From bc8d6b0a6d27b2baa568474c89fb1c3502256316 Mon Sep 17 00:00:00 2001 From: Tom Dendooven Date: Sat, 27 Jun 2026 14:30:25 +0200 Subject: [PATCH] Add optional gain reference to reconstruct (create_settings --gain_path) Frame series that are not pre-gain-corrected need a gain reference. WarpTools applies this at create_settings via --gain_path (+ optional --gain_flip_x/y, --gain_transpose). This threads an optional --gain argument (and orientation switches) from the CLI through reconstruct() into the frame-series create_settings call. Default None is fully backward-compatible: no gain arguments are emitted when unset. --- src/easymode/core/warp_wrapper.py | 15 +++++++++++++-- src/easymode/main.py | 10 +++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/src/easymode/core/warp_wrapper.py b/src/easymode/core/warp_wrapper.py index d47f0ff..cb8f003 100644 --- a/src/easymode/core/warp_wrapper.py +++ b/src/easymode/core/warp_wrapper.py @@ -123,10 +123,11 @@ def get_gpu_list(): except: return [] -def reconstruct(frames, mdocs, apix=None, dose=None, extension=None, tomo_apix=10.0, thickness=3000, shape=None, steps='1111111', halfmaps=True, force_align=False): +def reconstruct(frames, mdocs, apix=None, dose=None, extension=None, tomo_apix=10.0, thickness=3000, shape=None, steps='1111111', halfmaps=True, force_align=False, gain=None, gain_flip_x=False, gain_flip_y=False, gain_transpose=False): root = os.getcwd() frames_path = frames if os.path.exists(frames) else os.path.join(root, frames) mdoc_path = mdocs if os.path.exists(mdocs) else os.path.join(root, mdocs) + gain_path = (gain if os.path.exists(gain) else os.path.join(root, gain)) if gain else None extension = extension if extension is not None else find_extension(frames_path) extension = f'.{extension}' if not '.' in extension else extension @@ -140,6 +141,7 @@ def reconstruct(frames, mdocs, apix=None, dose=None, extension=None, tomo_apix=1 f'\ntomo apix: {tomo_apix}' f'\nthickness: {thickness}' f'\nshape: {shape if shape is not None else "auto"}' + f'\ngain: {gain_path if gain_path is not None else "none (frames assumed gain-corrected)"}' f'\nsteps: {steps}') if dose is None: @@ -153,7 +155,16 @@ def reconstruct(frames, mdocs, apix=None, dose=None, extension=None, tomo_apix=1 print(f'\n\033[96mCreating settings (frame series)\033[0m') - _run(f'WarpTools create_settings --folder_data {frames_path} --folder_processing warp_frameseries --output warp_frameseries.settings --extension "*{extension}" --angpix {apix} --exposure {dose}') + gain_args = '' + if gain_path is not None: + gain_args = f' --gain_path {gain_path}' + if gain_flip_x: + gain_args += ' --gain_flip_x' + if gain_flip_y: + gain_args += ' --gain_flip_y' + if gain_transpose: + gain_args += ' --gain_transpose' + _run(f'WarpTools create_settings --folder_data {frames_path} --folder_processing warp_frameseries --output warp_frameseries.settings --extension "*{extension}" --angpix {apix} --exposure {dose}{gain_args}') print(f'\n\033[96mCreating settings (tilt series)\033[0m') tomo_size = [int(f) for f in shape.split('x')] if shape is not None else find_shape(frames_path, extension) diff --git a/src/easymode/main.py b/src/easymode/main.py index c137fcc..8eb815b 100644 --- a/src/easymode/main.py +++ b/src/easymode/main.py @@ -61,6 +61,10 @@ def _parse_size(s): reconstruct.add_argument('--steps', type=str, default='11111111', help="8-character string indicating which processing steps to perform (default: '1111111'). Each character corresponds to a specific step: 1 to perform the step, 0 to skip it. The steps are: 1) Frame motion and CTF, 2) Importing tilt series, 3) Creating tilt stacks, 4) Tilt series alignment, 5) Import alignments, 6) Tilt series CTF, 7) Check handedness, 8) Reconstruct volumes.") reconstruct.add_argument('--no_halfmaps', dest='halfmaps', action='store_false', help="If set, do not generate half-maps during motion correction or tomogram reconstruction. This precludes most methods of denoising.") reconstruct.add_argument('--force_align', action='store_true', help="If set, force AreTomo3 alignment of tilt series even if alignment files are already present.") + reconstruct.add_argument('--gain', type=str, default=None, help="Path to a gain reference file. Applied during WarpTools create_settings (frame series). Leave empty when frames are already gain-corrected.") + reconstruct.add_argument('--gain_flip_x', action='store_true', help="Flip the gain reference along the X axis.") + reconstruct.add_argument('--gain_flip_y', action='store_true', help="Flip the gain reference along the Y axis.") + reconstruct.add_argument('--gain_transpose', action='store_true', help="Transpose the gain reference.") segment = subparsers.add_parser('segment', help='Segment data using pretrained easymode networks.') segment.add_argument( "features", metavar='FEATURE', nargs="+", type=str, help="One or more features to segment (e.g. 'ribosome membrane microtubule'). Use 'easymode list' to see available features.") @@ -275,7 +279,11 @@ def _parse_size(s): shape=args.shape, steps=args.steps, halfmaps=args.halfmaps, - force_align=args.force_align) + force_align=args.force_align, + gain=args.gain, + gain_flip_x=args.gain_flip_x, + gain_flip_y=args.gain_flip_y, + gain_transpose=args.gain_transpose) elif args.command == 'set': if args.cache_directory: if os.path.exists(args.cache_directory):