diff --git a/extensions/chaeo/params.py b/extensions/chaeo/params.py
index fccc4f0bae2943ed0ab63f2a1a298d15dda71e0f..ee4dabd6b579874dcce875689dd074082267b9e2 100644
--- a/extensions/chaeo/params.py
+++ b/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):
diff --git a/extensions/chaeo/tests/test_zstack.py b/extensions/chaeo/tests/test_zstack.py
index f472c7734bff93eedb880d1581e5bc84cf503c84..be14b104fbb17ece4bf228e532e9edcc7043c08a 100644
--- a/extensions/chaeo/tests/test_zstack.py
+++ b/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'],
diff --git a/extensions/chaeo/workflows.py b/extensions/chaeo/workflows.py
index 811aed5f6337d96a7a547e623657785ccc3fc00a..5e13b819ed9d84be41330f215eebcc0211f2e024 100644
--- a/extensions/chaeo/workflows.py
+++ b/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,
diff --git a/extensions/chaeo/zmask.py b/extensions/chaeo/zmask.py
index 77854e3e5f0b6f68394b9f4357c0179c713a0a27..3e16b8adccb8de1953334f1076521298a88c3c95 100644
--- a/extensions/chaeo/zmask.py
+++ b/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(