simplifying export calls

model_server/extensions/chaeo/params.py

@@ -36,7 +36,7 @@ class RoiSetExportParams(BaseModel):
     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
+    patch_masks: Union[PatchParams, None] = None
     annotated_zstacks: Union[AnnotatedZStackParams, None] = None
     object_classes: bool = False
     dataframe: bool = False

model_server/extensions/chaeo/products.py

@@ -60,9 +60,9 @@ def write_patch_to_file(where, fname, yxcz):
 
 def get_patch_masks(roiset, pad_to: int = 256) -> MonoPatchStack:
     patches = []
-    for ob in roiset.get_df().itertuples('Roi'):
-        patch = np.zeros((ob.ebb_h, ob.ebb_w, 1, 1), dtype='uint8')
-        patch[ob.relative_slice][:, :, 0, 0] = ob.mask * 255
+    for roi in roiset.get_df().itertuples():
+        patch = np.zeros((roi.ebb_h, roi.ebb_w, 1, 1), dtype='uint8')
+        patch[roi.relative_slice][:, :, 0, 0] = roi.mask * 255
 
         if pad_to:
             patch = pad(patch, pad_to)
@@ -72,15 +72,13 @@ def get_patch_masks(roiset, pad_to: int = 256) -> MonoPatchStack:
 
 
 def export_patch_masks(roiset, where: Path, pad_to: int = 256, prefix='mask') -> list:
-    patches_acc = roiset.get_patch_masks(pad_to=pad_to)
+    patches_acc = get_patch_masks(roiset, pad_to=pad_to)
 
     exported = []
-    for i in range(0, roiset.count):
-        mi = roiset.zmask_meta[i]
-        obj = mi['info']
+    for i, roi in enumerate(roiset.get_df().itertuples()):  # assumes index of patches_acc is same as dataframe
         patch = patches_acc.iat_yxcz(i)
         ext = 'png'
-        fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}.{ext}'
+        fname = f'{prefix}-la{roi.label:04d}-zi{roi.zi:04d}.{ext}'
         write_patch_to_file(where, fname, patch)
         exported.append(fname)
     return exported

model_server/extensions/chaeo/zmask.py

@@ -181,9 +181,6 @@ class RoiSet(object):
     def get_object_mask_by_class(self, class_id):
         return self.object_id_labels == class_id
 
-    def get_patch_masks(self, **kwargs) -> MonoPatchStack:
-        return get_patch_masks(self, **kwargs)
-
     def export_patch_masks(self, where, **kwargs) -> list:
         return export_patch_masks(self, where, **kwargs)
 
@@ -284,7 +281,7 @@ class RoiSet(object):
                 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':
-                self.export_patch_masks(subdir, prefix=pr)
+                self.export_patch_masks(subdir, prefix=pr, **params.patch_masks)
             if k == 'annotated_zstacks':
                 annotated = InMemoryDataAccessor(
                     draw_boxes_on_3d_image(raw_ch.data, self.zmask_meta, **kp)