Add optional gain reference to reconstruct (create_settings --gain_path)

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)

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):