diff --git a/jpeg_tomogram.py b/jpeg_tomogram.py index 4a88448..2da906f 100755 --- a/jpeg_tomogram.py +++ b/jpeg_tomogram.py @@ -200,15 +200,16 @@ def mrc_to_jpeg_stack(mrc_filename, jpeg_stack_filename, quality, cores=None, ve print_success(f"{jpeg_stack_filename} successfully packed!") return jpeg_stack_filename -def jpeg_stack_to_mrc(jpeg_stack_filename, mrc_filename, cores=None, verbose=False, quiet=False): +def jpeg_stack_to_mrc(jpeg_stack_filename, mrc_filename, cores=None, verbose=False, output_format='mrc', quiet=False): """ - Convert a JPEG stack to an MRC file. + Convert a JPEG stack to an MRC file or a directory of JPEGS. :param str jpeg_stack_filename: The input JPEG stack filename - :param str mrc_filename: The output MRC filename + :param str mrc_filename: The output MRC filename or directory :param int cores: Number of CPU cores to use (default: None, uses all available cores) :param bool verbose: Whether to print verbose output :param bool quiet: Whether to suppress all output and progress bars + :param str output_format: Output format, either 'mrc' or 'dir' """ if verbose and not quiet: print(f"Input JPEG stack: {jpeg_stack_filename}") @@ -217,19 +218,23 @@ def jpeg_stack_to_mrc(jpeg_stack_filename, mrc_filename, cores=None, verbose=Fal if cores is None: cores = cpu_count() - # Find the header file - base_name = os.path.splitext(jpeg_stack_filename)[0] - header_pattern = f"{base_name}*_header.npy" - matching_headers = list(Path('.').glob(header_pattern)) - - if matching_headers: - header_file = str(matching_headers[0]) - # Extract the original filename from the header filename - original_filename = header_file.rsplit('_JPG', 1)[0] + '.mrc' - mrc_filename = original_filename + if output_format == 'dir': + output_dir = Path(mrc_filename).with_suffix('') + output_dir.mkdir(parents=True, exist_ok=True) else: - # If no header file is found, use the input filename without the .jpgs extension - mrc_filename = base_name + '.mrc' + # Find the header file + base_name = os.path.splitext(jpeg_stack_filename)[0] + header_pattern = f"{base_name}*_header.npy" + matching_headers = list(Path('.').glob(header_pattern)) + + if matching_headers: + header_file = str(matching_headers[0]) + # Extract the original filename from the header filename + original_filename = header_file.rsplit('_JPG', 1)[0] + '.mrc' + mrc_filename = original_filename + else: + # If no header file is found, use the input filename without the .jpgs extension + mrc_filename = base_name + '.mrc' if verbose: print(f"Output MRC filename (after processing): {mrc_filename}") @@ -261,18 +266,24 @@ def jpeg_stack_to_mrc(jpeg_stack_filename, mrc_filename, cores=None, verbose=Fal data = np.array([np.array(img) for img in (images if quiet else tqdm(images, desc="Converting", unit="slice"))]) - header = read_header(jpeg_stack_filename) + if output_format == 'dir': + for i, img in enumerate(tqdm(images, desc="Saving JPGs", unit="slice")): + img.save(output_dir / f"{i:04d}.jpg", "JPEG") + print_success(f"JPGs successfully unpacked to {output_dir}!") - with mrcfile.new(mrc_filename, overwrite=True) as mrc: - include_fields = ['nx', 'ny', 'nz', 'mode', 'nxstart', 'nystart', 'nzstart', 'mx', 'my', 'mz', 'xlen', 'ylen', 'zlen', 'alpha', 'beta', 'gamma', 'mapc', 'mapr', 'maps', 'amin', 'amax', 'amean', 'ispg', 'extra', 'xorigin', 'yorigin', 'zorigin', 'map', 'machst', 'rms', 'nlabels', 'cella'] - for field in include_fields: - try: - mrc.header[field] = header[field] - except: - pass - mrc.set_data((data - 128).astype(np.int8)) - if not quiet: - print_success(f"{mrc_filename} successfully unpacked!") + else: + header = read_header(jpeg_stack_filename) + + with mrcfile.new(mrc_filename, overwrite=True) as mrc: + include_fields = ['nx', 'ny', 'nz', 'mode', 'nxstart', 'nystart', 'nzstart', 'mx', 'my', 'mz', 'xlen', 'ylen', 'zlen', 'alpha', 'beta', 'gamma', 'mapc', 'mapr', 'maps', 'amin', 'amax', 'amean', 'ispg', 'extra', 'xorigin', 'yorigin', 'zorigin', 'map', 'machst', 'rms', 'nlabels', 'cella'] + for field in include_fields: + try: + mrc.header[field] = header[field] + except: + pass + mrc.set_data((data - 128).astype(np.int8)) + if not quiet: + print_success(f"{mrc_filename} successfully unpacked!") return mrc_filename def report_compression_ratio(input_path, output_files): @@ -318,6 +329,33 @@ def unpack_file(args): """ return jpeg_stack_to_mrc(*args) +def collect_input_paths(input_paths, mode): + """ + Collects and returns a list of valid input paths based on the specified mode. + + :param input_paths: List of input paths (files or directories) + :param mode: Operation mode, either 'pack' or 'unpack' + :return: List of paths to process + """ + valid_extensions = {'.mrc', '.rec'} if mode == 'pack' else {'.jpgs'} + collected_paths = [] + + for path in input_paths: + path_obj = Path(path) + + if path_obj.is_dir(): + collected_paths.extend(path_obj.glob(f"*{ext}") for ext in valid_extensions) + elif path_obj.is_file() and path_obj.suffix in valid_extensions: + collected_paths.append(path_obj) + else: + # Handle wildcard patterns (e.g., '*', '?', or '[') + if any(char in path for char in '*?[]'): + collected_paths.extend( + f for f in Path().glob(path) if f.suffix in valid_extensions + ) + + return collected_paths + def main(): """ Main function to handle command-line arguments and execute packing or unpacking operations. @@ -328,6 +366,7 @@ def main(): parser.add_argument('mode', choices=['pack', 'unpack'], help='Whether to pack an MRC file into a JPEG stack or unpack a JPEG stack into an MRC file') parser.add_argument('input_paths', nargs='+', help='The input file(s) or directory of files (.mrc and .rec supported). Can be a list or a wildcard expression.') parser.add_argument('-o', '--output_path', help='The output file or directory (default: same as input path)') + parser.add_argument('-f', '--output_format', choices=['mrc', 'dir'], default='mrc', help="Specify the output format when unpacking. 'mrc' for MRC file, 'dir' for a directory of JPGs (default: mrc)") parser.add_argument('-e', '--external_viewer', help='External program to open the unpacked MRC file(s) (e.g. 3dmod)') parser.add_argument('-q', '--quality', type=validate_quality, default=80, help='The quality of the JPEG images in the stack (1-100). Note: values above 95 should be avoided (default: 80)') parser.add_argument('-c', '--cores', type=int, default=None, help='Number of CPU cores to use (default: all)') @@ -336,30 +375,7 @@ def main(): parser.add_argument("-v", "--version", action="version", help="Show version number and exit", version=f"JPEG Tomogram v{__version__}") args = parser.parse_args() - # If quiet mode is enabled, suppress all output - if args.quiet: - sys.stdout = open(os.devnull, 'w') - sys.stderr = open(os.devnull, 'w') - - input_paths = [] - for path in args.input_paths: - if Path(path).is_dir(): - if args.mode == 'pack': - input_paths.extend(Path(path).glob('*.mrc')) - input_paths.extend(Path(path).glob('*.rec')) - elif args.mode == 'unpack': - input_paths.extend(Path(path).glob('*.jpgs')) - else: - # Handle wildcard patterns (e.g., '*') - if '*' in path or '?' in path or '[' in path: - # Filter to include only .mrc and .rec files - for f in Path().glob(path): - if f.suffix in {'.mrc', '.rec'}: - input_paths.append(f) - else: - # Explicitly check the file extension - if Path(path).suffix in {'.mrc', '.rec'}: - input_paths.append(Path(path)) + input_paths = collect_input_paths(args.input_paths, args.mode) if not input_paths: print_error(f"Error: No matching files found for input path(s): {args.input_paths}") @@ -394,7 +410,8 @@ def main(): str(f), str((output_path or f.parent) / (f.stem.replace(f'_JPG{args.quality}', f'_fromJPG{args.quality}'))), 1, - args.verbose + args.verbose, + args.output_format ) for f in input_paths ]