Export patch masks

extensions/chaeo/products.py

@@ -42,6 +42,26 @@ def _write_patch_to_file(where, fname, data):
     else:
         raise Exception(f'Unsupported file extension: {ext}')
 
+def export_patch_masks_from_zstack(
+    where: Path,
+    zmask_meta: list,
+    pad_to: int = 256,
+    prefix='mask',
+):
+    exported = []
+    for mi in zmask_meta:
+        obj = mi['info']
+        mask = np.expand_dims(mi['mask'], (2, 3))
+
+        if pad_to:
+            mask = pad(mask, pad_to)
+
+        ext = 'png'
+        fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}.{ext}'
+        mask8bit = 255 * mask.astype('uint8')
+        _write_patch_to_file(where, fname, mask8bit)
+        exported.append(fname)
+    return exported
 
 def export_patches_from_zstack(
         where: Path,

extensions/chaeo/workflows.py

@@ -3,7 +3,7 @@ from typing import Dict
 
 from extensions.ilastik.models import IlastikPixelClassifierModel
 from extensions.chaeo.annotators import draw_boxes_on_3d_image
-from extensions.chaeo.products import export_patches_from_zstack
+from extensions.chaeo.products import export_patches_from_zstack, export_patch_masks_from_zstack
 from extensions.chaeo.zmask import build_zmask_from_object_mask, project_stack_from_focal_points
 from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
 from model_server.workflows import Timer
@@ -20,6 +20,11 @@ def export_patches_from_multichannel_zstack(
         mask_type: str = 'boxes',
         zmask_filters: Dict = None,
         zmask_expand_box_by: int = None,
+        export_pixel_probabilities=True,
+        export_2d_patches=True,
+        export_3d_patches=True,
+        export_annotated_zstack=True,
+        export_patch_masks=True,
 ) -> Dict:
 
     ti = Timer()
@@ -38,11 +43,12 @@ def export_patches_from_multichannel_zstack(
     pxmap, _ = px_model.infer(mip)
     ti.click('infer_pixel_probability')
 
-    write_accessor_data_to_file(
-        Path(where_output) / 'pixel_probabilities' / (fstem + '.tif'),
-        pxmap
-    )
-    ti.click('export_pixel_probability')
+    if export_pixel_probabilities:
+        write_accessor_data_to_file(
+            Path(where_output) / 'pixel_probabilities' / (fstem + '.tif'),
+            pxmap
+        )
+        ti.click('export_pixel_probability')
 
     obmask = InMemoryDataAccessor(
         pxmap.data > pxmap_threshold
@@ -60,51 +66,59 @@ def export_patches_from_multichannel_zstack(
     zmask_acc = InMemoryDataAccessor(zmask)
     ti.click('generate_zmasks')
 
-    files = export_patches_from_zstack(
-        Path(where_output) / '3d_patches',
-        stack.get_one_channel_data(patches_channel),
-        zmask_meta,
-        prefix=fstem,
-        draw_bounding_box=False,
-        rescale_clip=0.0,
-        make_3d=True,
-    )
-    ti.click('export_patches')
+    if export_3d_patches:
+        files = export_patches_from_zstack(
+            Path(where_output) / '3d_patches',
+            stack.get_one_channel_data(patches_channel),
+            zmask_meta,
+            prefix=fstem,
+            draw_bounding_box=False,
+            rescale_clip=0.0,
+            make_3d=True,
+        )
+        ti.click('export_3d_patches')
 
-    files = export_patches_from_zstack(
-        Path(where_output) / '2d_patches',
-        stack.get_one_channel_data(patches_channel),
-        zmask_meta,
-        prefix=fstem,
-        draw_bounding_box=False,
-        rescale_clip=0.0,
-        make_3d=False,
-        focus_metric='max_sobel',
-    )
-    ti.click('export_patches')
+    if export_2d_patches:
+        files = export_patches_from_zstack(
+            Path(where_output) / '2d_patches',
+            stack.get_one_channel_data(patches_channel),
+            zmask_meta,
+            prefix=fstem,
+            draw_bounding_box=False,
+            rescale_clip=0.0,
+            make_3d=False,
+            focus_metric='max_sobel',
+        )
+        ti.click('export_2d_patches')
 
-    # export annotated zstack
-    annotated = InMemoryDataAccessor(
-        draw_boxes_on_3d_image(
-            stack.get_one_channel_data(patches_channel).data,
-            zmask_meta
+    if export_patch_masks:
+        files = export_patch_masks_from_zstack(
+            Path(where_output) / 'patch_masks',
+            zmask_meta,
         )
-    )
-    write_accessor_data_to_file(
-        Path(where_output) / 'annotated_zstacks' / (fstem + '.tif'),
-        annotated
-    )
-    ti.click('export_annotated_zstack')
 
-    # # generate multichannel projection from label centroids
-    # dff = df[df['keeper']]
-    # interm['projected'] = project_stack_from_focal_points(
-    #     dff['centroid-0'].to_numpy(),
-    #     dff['centroid-1'].to_numpy(),
-    #     dff['zi'].to_numpy(),
-    #     stack,
-    #     degree=4,
-    # )
+    if export_annotated_zstack:
+        annotated = InMemoryDataAccessor(
+            draw_boxes_on_3d_image(
+                stack.get_one_channel_data(patches_channel).data,
+                zmask_meta
+            )
+        )
+        write_accessor_data_to_file(
+            Path(where_output) / 'annotated_zstacks' / (fstem + '.tif'),
+            annotated
+        )
+        ti.click('export_annotated_zstack')
+
+    # generate multichannel projection from label centroids
+    dff = df[df['keeper']]
+    interm['projected'] = project_stack_from_focal_points(
+        dff['centroid-0'].to_numpy(),
+        dff['centroid-1'].to_numpy(),
+        dff['zi'].to_numpy(),
+        stack,
+        degree=4,
+    )
 
     return {
         'pixel_model_id': px_model.model_id,