Do not perform sum of channels in InputFile

Make channel arguments uniform
lens-biophotonics · Jun 9, 2022 · 348c7ba · 348c7ba
1 parent d5f00f8
commit 348c7ba
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 19 deletions.
diff --git a/zetastitcher/align/__main__.py b/zetastitcher/align/__main__.py
@@ -54,7 +54,7 @@ def parse_args():
 
     parser.add_argument('input_folder', help='input folder')
     parser.add_argument('-o', type=str, default='stitch.yml', dest='output_file', help='output file')
-    parser.add_argument('-c', type=str, default='s', dest='channel', choices=['r', 'g', 'b', 's'], help='color channel')
+    parser.add_argument('-c', '--ch', type=int, dest='channel', help='color channel')
     parser.add_argument('-j', type=int, dest='n_of_workers',
                         help='number of parallel jobs (defaults to number of system cores)')
     parser.add_argument('-r', action='store_true', dest='recursive', help='recursively look for files')
@@ -120,15 +120,6 @@ def parse_args():
         setattr(args, 'overlap_h', args.overlap)
         setattr(args, 'overlap_v', args.overlap)
 
-    channels = {
-        's': -2,  # sum
-        'r': 0,
-        'g': 1,
-        'b': 2
-    }
-
-    args.channel = channels[args.channel]
-
     args.max_dx = int(round(args.max_dx / args.px_size_xy))
     args.max_dy = int(round(args.max_dy / args.px_size_xy))
     args.max_dz = int(round(args.max_dz / args.px_size_z))
@@ -157,18 +148,25 @@ def worker(item, overlap_dict, channel, max_dz, max_dy, max_dx):
     a = InputFile(aname)
     b = InputFile(bname)
 
-    a.channel = channel
-    b.channel = channel
-
     z_min = z_frame - max_dz
     z_max = z_frame + max_dz + 1
 
     aslice = a.zslice(z_min, z_max, copy=True)
+    if a.nchannels > 1:
+        if channel is not None:
+            aslice = aslice[:, channel]
+        else:
+            aslice = np.sum(aslice.astype(np.float32), axis=1)
     if axis == 2:
         aslice = np.rot90(aslice, axes=(-1, -2))
     aslice = aslice[..., -(overlap + max_dy):, :]
 
     bframe = b.zslice_idx(z_frame, copy=True)
+    if b.nchannels > 1:
+        if channel is not None:
+            bframe = bframe[:, channel]
+        else:
+            bframe = np.sum(bframe.astype(np.float32), axis=1)
     if axis == 2:
         bframe = np.rot90(bframe, axes=(-1, -2))
     bframe = bframe[..., :overlap - max_dy, :]

diff --git a/zetastitcher/fuse/__main__.py b/zetastitcher/fuse/__main__.py
@@ -55,7 +55,7 @@ def parse_args():
     group.add_argument('-d', dest='debug', action='store_true',
                        help='overlay debug info')
 
-    group.add_argument('--ch', type=int, dest='channel', help='channel')
+    group.add_argument('-c', '--ch', type=int, dest='channel', help='channel')
 
     group.add_argument('--downsample-xy', metavar='S', type=int, required=False,
                        help='downsample XY plane by factor S')

diff --git a/zetastitcher/io/inputfile.py b/zetastitcher/io/inputfile.py
@@ -23,7 +23,7 @@ def __init__(self, file_path=None):
         super().__init__()
         self.file_path = file_path
         self.wrapper = None
-        self._channel = -1
+        self._channel = None
         self.squeeze = True
 
         self.nfrms = None
@@ -103,7 +103,7 @@ def shape(self):
             the `channels` dimension is squeezed.
         """
         s = [self.nfrms, self.nchannels, self.ysize, self.xsize]
-        if self.nchannels == 1 or self.channel != -1:
+        if self.nchannels == 1 or self.channel is not None:
             del s[1]
         return tuple(s)
 
@@ -223,9 +223,7 @@ def zslice(self, arg1, arg2=None, step=None, dtype=None, copy=True):
                 a[z] = self.wrapper.frame(i)
                 z += 1
 
-        if self.channel == -2:
-            a = np.sum(a, axis=-1)
-        elif self.channel != -1:
+        if self.channel is not None:
             a = a[..., self.channel]
         elif self.nchannels > 1 and a.ndim >= 3:
             a = np.moveaxis(a, -1, -3)