Combined all export products into a single call, which routes individual...

Combined all export products into a single call, which routes individual products' export methods and parameters

extensions/chaeo/params.py

@@ -17,15 +17,17 @@ class PatchParams(BaseModel):
 class AnnotatedZStackParams(BaseModel):
     draw_label: bool = False
 
+class RoiSetExportMetaParams(BaseModel):
+    expand_box_by: List[int] = [128, 0]
 
 class RoiSetExportParams(BaseModel):
-    expand_box_by: List[int] = [128, 0]
+    meta: RoiSetExportMetaParams = RoiSetExportMetaParams()
     pixel_probabilities: bool = False
     patches_3d: Union[PatchParams, None] = None
     patches_2d_for_annotation: Union[PatchParams, None] = None
     patches_2d_for_training: Union[PatchParams, None] = None
     patch_masks: bool = False
-    annotated_z_stack: Union[AnnotatedZStackParams, None] = None
+    annotated_zstacks: Union[AnnotatedZStackParams, None] = None
 
 
 class RoiFilterRange(BaseModel):

extensions/chaeo/tests/test_zstack.py

@@ -223,7 +223,7 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
                 'draw_mask': False,
             },
             'patch_masks': True,
-            'annotated_z_stack': {}
+            'annotated_zstacks': {}
         })
         infer_object_map_from_zstack(
             multichannel_zstack['path'],

extensions/chaeo/workflows.py

@@ -76,7 +76,7 @@ def infer_object_map_from_zstack(
         stack,
         mask_type=zmask_type,
         filters=zmask_filters,
-        expand_box_by=exports.expand_box_by,
+        expand_box_by=exports.meta.expand_box_by,
     )
     ti.click('generate_zmasks')
 
@@ -90,48 +90,7 @@ def infer_object_map_from_zstack(
     )
     ti.click('export_object_classes')
 
-    if exports.patches_3d:
-            rois.export_3d_patches(
-                Path(output_folder_path) / '3d_patches',
-                fstem,
-                patches_channel,
-                exports.patches_3d
-            )
-    ti.click('export_3d_patches')
-
-    if exports.patches_2d_for_annotation:
-        rois.export_2d_patches_for_annotation(
-            Path(output_folder_path) / '2d_patches_annotation',
-            fstem,
-            patches_channel,
-            exports.patches_2d_for_annotation
-        )
-    ti.click('export_2d_patches_for_annotation')
-
-    if exports.patches_2d_for_training:
-        rois.export_2d_patches_for_training(
-            Path(output_folder_path) / '2d_patches_training',
-            fstem,
-            patches_channel,
-            exports.patches_2d_for_training
-        )
-    ti.click('export_2d_patches_for_training')
-
-    if exports.patch_masks:
-        rois.export_patch_masks(
-            Path(output_folder_path) / 'patch_masks',
-            fstem,
-            patches_channel,
-        )
-
-    if exports.annotated_z_stack:
-        rois.export_annotated_zstack(
-            Path(output_folder_path) / 'patch_masks',
-            fstem,
-            patches_channel,
-            exports.annotated_z_stack
-        )
-    ti.click('export_annotated_zstack')
+    rois.run_exports(Path(output_folder_path), patches_channel, exports)
 
     return {
         'timer_results': ti.events,

extensions/chaeo/zmask.py

@@ -40,72 +40,6 @@ class RoiSet(object):
     def get_argmax(self):
         return self.interm.argmax
 
-    def export_3d_patches(self, where, prefix, channel, params: PatchParams):
-        if not self.count:
-            return
-        files = export_patches_from_zstack(
-            where,
-            self.acc_raw.get_one_channel_data(channel),
-            self.zmask_meta,
-            prefix=prefix,
-            make_3d=True,
-            **params.__dict__,
-        )
-
-    def export_2d_patches_for_training(self, where, prefix, channel, params: PatchParams):
-        if not self.count:
-            return
-        files = export_multichannel_patches_from_zstack(
-            where,
-            self.acc_raw,
-            self.zmask_meta,
-            ch_white=channel,
-            prefix=prefix,
-            make_3d=False,
-            **params.__dict__,
-        )
-        df_patches = pd.DataFrame(files)
-        self.df = pd.merge(self.df, df_patches, left_index=True, right_on='df_index').drop(columns='df_index')
-        self.df['patch_id'] = self.df.apply(lambda _: uuid4(), axis=1)
-
-    def export_2d_patches_for_annotation(self, where, prefix, channel, params: PatchParams):
-        if not self.count:
-            return
-        files = export_multichannel_patches_from_zstack(
-            where,
-            self.acc_raw,
-            self.zmask_meta,
-            prefix=prefix,
-            make_3d=False,
-            ch_white=channel,
-            bounding_box_channel=1,
-            bounding_box_linewidth=2,
-            **params.__dict__,
-        )
-
-    def export_patch_masks(self, where, prefix, channel):
-        if not self.count:
-            return
-        files = export_patch_masks_from_zstack(
-            where,
-            self.acc_raw.get_one_channel_data(channel),
-            self.zmask_meta,
-            prefix=prefix,
-        )
-
-    def export_annotated_zstack(self, where, prefix, channel, params: AnnotatedZStackParams):
-        annotated = InMemoryDataAccessor(
-            draw_boxes_on_3d_image(
-                self.acc_raw.get_one_channel_data(channel).data,
-                self.zmask_meta,
-                **params.__dict__,
-            )
-        )
-        write_accessor_data_to_file(
-            where / 'annotated_zstacks' / (prefix + '.tif'),
-            annotated
-        )
-
     def get_multichannel_projection(self):
         dff = self.df[self.df['keeper']]
         if self.count:
@@ -164,15 +98,41 @@ class RoiSet(object):
     def get_object_map(self, filters: RoiFilter):
         pass
 
-    def run_exports(self, where, channel, params):
-        names = [
-            'pixel_probabilities',
-            'patches_3d',
-            'patches_2d_for_annotation',
-            'patches_2d_for_training',
-            'patch_masks',
-            'annotated_z_stack',
-        ]
+    def run_exports(self, where, channel, params: RoiSetExportParams):
+        if not self.count:
+            return
+        raw_ch = self.acc_raw.get_one_channel_data(channel)
+        for k in params.dict().keys():
+            subdir = where / k
+            kp = params.dict()[k]
+            if k == 'meta' or kp is None:
+                continue
+            if k == '3d_patches':
+                files = export_patches_from_zstack(
+                    subdir, raw_ch, self.zmask_meta, prefix=k, make_3d=True, **kp
+                )
+            if k == 'patches_2d_for_annotation':
+                files = export_multichannel_patches_from_zstack(
+                    subdir, self.acc_raw, self.zmask_meta, prefix=k, make_3d=False, ch_white=channel,
+                    bounding_box_channel=1, bounding_box_linewidth=2, **kp,
+                )
+            if k == 'patches_2d_for_training':
+                files = export_multichannel_patches_from_zstack(
+                    subdir, self.acc_raw, self.zmask_meta, ch_white=channel, prefix=k, make_3d=False, **kp
+                )
+                df_patches = pd.DataFrame(files)
+                self.df = pd.merge(self.df, df_patches, left_index=True, right_on='df_index').drop(columns='df_index')
+                self.df['patch_id'] = self.df.apply(lambda _: uuid4(), axis=1)
+            if k == 'patch_masks':
+                export_patch_masks_from_zstack(
+                    subdir, raw_ch, self.zmask_meta, prefix=k,
+                )
+            if k == 'annotated_zstacks':
+                annotated = InMemoryDataAccessor(
+                    draw_boxes_on_3d_image(raw_ch.data, self.zmask_meta, **kp)
+                )
+                write_accessor_data_to_file(subdir / (k + '.tif'), annotated)
+
 
 
 def build_zmask_from_object_mask(