Delegated export of patch products to ZMaskObjectTable class; consolidated...

Delegated export of patch products to ZMaskObjectTable class; consolidated both batch runner and API-facing workflows into infer_object_map_from_zstack(); consolidated parameterization of export products into pydantic model

extensions/chaeo/examples/export_patch_focus_metrics.py

@@ -9,7 +9,7 @@ from skimage.filters import gaussian, sobel
 
 from extensions.ilastik.models import IlastikPixelClassifierModel
 from extensions.chaeo.products import export_3d_patches_with_focus_metrics, export_patches_from_zstack
-from extensions.chaeo.zmask import build_zmask_from_object_mask
+from extensions.chaeo.zmask import ZMaskObjectTable
 from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
 from model_server.workflows import Timer
 
@@ -51,20 +51,27 @@ def export_patch_focus_metrics_from_multichannel_zstack(
     ti.click('threshold_pixel_mask')
 
     # make zmask
-    zmask, zmask_meta, df, interm = build_zmask_from_object_mask(
+    # zmask, zmask_meta, df, interm = build_zmask_from_object_mask(
+    #     obmask.get_one_channel_data(pixel_class),
+    #     stack.get_one_channel_data(zmask_channel),
+    #     mask_type=mask_type,
+    #     filters=zmask_filters,
+    #     expand_box_by=zmask_expand_box_by,
+    # )
+    obj_table = ZMaskObjectTable(
         obmask.get_one_channel_data(pixel_class),
         stack.get_one_channel_data(zmask_channel),
         mask_type=mask_type,
         filters=zmask_filters,
         expand_box_by=zmask_expand_box_by,
     )
-    zmask_acc = InMemoryDataAccessor(zmask)
+    zmask_acc = InMemoryDataAccessor(obj_table.zmask)
     ti.click('generate_zmasks')
 
     files = export_3d_patches_with_focus_metrics(
         Path(where_output) / '3d_patches',
         stack.get_one_channel_data(patches_channel),
-        zmask_meta,
+        obj_table.zmask_meta,
         prefix=fstem,
         rescale_clip=0.0,
         make_3d=True,
@@ -75,7 +82,7 @@ def export_patch_focus_metrics_from_multichannel_zstack(
     files = export_patches_from_zstack(
         Path(where_output) / '2d_patches',
         stack.get_one_channel_data(patches_channel),
-        zmask_meta,
+        obj_table.zmask_meta,
         prefix=fstem,
         draw_bounding_box=True,
         rescale_clip=0.0,
@@ -88,11 +95,11 @@ def export_patch_focus_metrics_from_multichannel_zstack(
     return {
         'pixel_model_id': px_model.model_id,
         'input_filepath': input_zstack_path,
-        'number_of_objects': len(zmask_meta),
+        'number_of_objects': len(obj_table.zmask_meta),
         'success': True,
         'timer_results': ti.events,
         'dataframe': df,
-        'interm': interm,
+        'interm': obj_table.interm,
     }
 
 if __name__ == '__main__':

extensions/chaeo/params.py

+from typing import Dict, List, Union
+
+from pydantic import BaseModel
+
+class PatchParams(BaseModel):
+    draw_bounding_box: bool = False
+    draw_contour: bool = False
+    draw_mask: bool = False
+    rescale_clip: float = 0.001
+    focus_metric: str = 'max_sobel'
+    rgb_overlay_channels: List[int] = (None, None, None),
+    rgb_overlay_weights: List[float] = (1.0, 1.0, 1.0)
+
+class AnnotatedZStackParams(BaseModel):
+    draw_label: bool = False
+
+class ZMaskExportParams(BaseModel):
+    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
+

extensions/chaeo/workflows.py

@@ -4,6 +4,7 @@ from uuid import uuid4
 
 import numpy as np
 import pandas as pd
+
 from skimage.measure import label, regionprops_table
 from skimage.morphology import dilation
 from sklearn.model_selection import train_test_split
@@ -11,9 +12,10 @@ from sklearn.model_selection import train_test_split
 from extensions.chaeo.accessors import MonoPatchStack
 from extensions.chaeo.annotators import draw_boxes_on_3d_image
 from extensions.chaeo.models import PatchStackObjectClassifier
+from extensions.chaeo.params import ZMaskExportParams
 from extensions.chaeo.process import mask_largest_object
 from extensions.chaeo.products import export_patches_from_zstack, export_patch_masks_from_zstack, export_multichannel_patches_from_zstack, get_patches_from_zmask_meta, get_patch_masks_from_zmask_meta
-from extensions.chaeo.zmask import build_zmask_from_object_mask, project_stack_from_focal_points
+from extensions.chaeo.zmask import project_stack_from_focal_points, ZMaskObjectTable
 from extensions.ilastik.models import IlastikPixelClassifierModel
 
 from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
@@ -21,18 +23,232 @@ from model_server.models import Model
 from model_server.process import rescale
 from model_server.workflows import Timer
 
-def get_zmask_meta(
-    input_file_path: str,
-    ilastik_pixel_classifier: IlastikPixelClassifierModel,
-    segmentation_channel: int,
-    pxmap_threshold: float,
-    pxmap_foreground_channel: int = 0,
-    zmask_zindex: int = None,
-    zmask_clip: int = None,
-    zmask_filters: Dict = None,
-    zmask_type: str = 'boxes',
-    **kwargs,
-) -> tuple:
+# def get_zmask_meta(
+#     input_file_path: str,
+#     ilastik_pixel_classifier: IlastikPixelClassifierModel,
+#     segmentation_channel: int,
+#     pxmap_threshold: float,
+#     pxmap_foreground_channel: int = 0,
+#     zmask_zindex: int = None,
+#     zmask_clip: int = None,
+#     zmask_filters: Dict = None,
+#     zmask_type: str = 'boxes',
+#     **kwargs,
+# ) -> tuple:
+#     ti = Timer()
+#     stack = generate_file_accessor(Path(input_file_path))
+#     fstem = Path(input_file_path).stem
+#     ti.click('file_input')
+#
+#     # MIP if no zmask z-index is given, then classify pixels
+#     if isinstance(zmask_zindex, int):
+#         assert 0 < zmask_zindex < stack.nz
+#         zmask_data = stack.get_one_channel_data(channel=segmentation_channel).data[:, :, :, zmask_zindex]
+#     else:
+#         zmask_data = stack.get_one_channel_data(channel=segmentation_channel).data.max(axis=-1, keepdims=True)
+#     if zmask_clip:
+#         zmask_data = rescale(zmask_data, zmask_clip)
+#     mip = InMemoryDataAccessor(
+#         zmask_data,
+#     )
+#     pxmap, _ = ilastik_pixel_classifier.infer(mip)
+#     ti.click('infer_pixel_probability')
+#
+#     obmask = InMemoryDataAccessor(
+#         pxmap.data > pxmap_threshold
+#     )
+#     ti.click('threshold_pixel_mask')
+#
+#     # make zmask
+#     obj_table = ZMaskObjectTable(
+#         obmask.get_one_channel_data(pxmap_foreground_channel),
+#         stack.get_one_channel_data(segmentation_channel),
+#         mask_type=zmask_type,
+#         filters=zmask_filters,
+#         expand_box_by=kwargs['zmask_expand_box_by'],
+#     )
+#     ti.click('generate_zmasks')
+#
+#     # record pixel scale
+#     obj_table.df['pixel_scale_in_micrometers'] = float(stack.pixel_scale_in_micrometers.get('X'))
+#
+#     return ti, stack, fstem, obmask, pxmap, obj_table
+
+
+# # called by batch runners
+# def export_patches_from_multichannel_zstack(
+#         input_file_path: str,
+#         output_folder_path: str,
+#         models: List[Model],
+#         pxmap_threshold: float,
+#         pxmap_foreground_channel: int,
+#         segmentation_channel: int,
+#         patches_channel: int,
+#         zmask_zindex: int = None,  # None for MIP,
+#         zmask_clip: int = None,
+#         zmask_type: str = 'boxes',
+#         zmask_filters: Dict = None,
+#         zmask_expand_box_by: int = None,
+#         export_pixel_probabilities=True,
+#         export_2d_patches_for_training=True,
+#         export_2d_patches_for_annotation=True,
+#         draw_bounding_box_on_2d_patch=True,
+#         draw_contour_on_2d_patch=False,
+#         draw_mask_on_2d_patch=False,
+#         export_3d_patches=True,
+#         export_annotated_zstack=True,
+#         draw_label_on_zstack=False,
+#         export_patch_masks=True,
+#         rgb_overlay_channels=(None, None, None),
+#         rgb_overlay_weights=(1.0, 1.0, 1.0),
+# ) -> Dict:
+#     pixel_classifier = models[0]
+#
+#     # ti, stack, fstem, obmask, pxmap, obj_table = get_zmask_meta(
+#     #     input_file_path,
+#     #     pixel_classifier,
+#     #     segmentation_channel,
+#     #     pxmap_threshold,
+#     #     pxmap_foreground_channel=pxmap_foreground_channel,
+#     #     zmask_zindex=zmask_zindex,
+#     #     zmask_clip=zmask_clip,
+#     #     zmask_expand_box_by=zmask_expand_box_by,
+#     #     zmask_filters=zmask_filters,
+#     #     zmask_type=zmask_type,
+#     # )
+#
+#     # obj_table = ZMaskObjectTable(
+#     #     obmask.get_one_channel_data(pxmap_foreground_channel),
+#     #     stack.get_one_channel_data(segmentation_channel),
+#     #     mask_type=zmask_type,
+#     #     filters=zmask_filters,
+#     #     expand_box_by=kwargs['zmask_expand_box_by'],
+#     # )
+#
+#     if export_pixel_probabilities:
+#         write_accessor_data_to_file(
+#             Path(output_folder_path) / 'pixel_probabilities' / (fstem + '.tif'),
+#             pxmap
+#         )
+#         ti.click('export_pixel_probability')
+#
+#     if export_3d_patches and len(zmask_meta) > 0:
+#         files = export_patches_from_zstack(
+#             Path(output_folder_path) / '3d_patches',
+#             stack.get_one_channel_data(patches_channel),
+#             zmask_meta,
+#             prefix=fstem,
+#             draw_bounding_box=False,
+#             rescale_clip=0.001,
+#             make_3d=True,
+#         )
+#         ti.click('export_3d_patches')
+#
+#     if export_2d_patches_for_annotation and len(zmask_meta) > 0:
+#         files = export_multichannel_patches_from_zstack(
+#             Path(output_folder_path) / '2d_patches_annotation',
+#             stack,
+#             zmask_meta,
+#             prefix=fstem,
+#             rescale_clip=0.001,
+#             make_3d=False,
+#             focus_metric='max_sobel',
+#             ch_white=patches_channel,
+#             ch_rgb_overlay=rgb_overlay_channels,
+#             draw_bounding_box=draw_bounding_box_on_2d_patch,
+#             bounding_box_channel=1,
+#             bounding_box_linewidth=2,
+#             draw_contour=draw_contour_on_2d_patch,
+#             draw_mask=draw_mask_on_2d_patch,
+#             overlay_gain=rgb_overlay_weights,
+#         )
+#         df_patches = pd.DataFrame(files)
+#         ti.click('export_2d_patches')
+#         # associate 2d patches, dropping labeled objects that were not exported as patches
+#         df = pd.merge(df, df_patches, left_index=True, right_on='df_index').drop(columns='df_index')
+#         # prepopulate patch UUID
+#         df['patch_id'] = df.apply(lambda _: uuid4(), axis=1)
+#
+#     if export_2d_patches_for_training and len(zmask_meta) > 0:
+#         files = export_multichannel_patches_from_zstack(
+#             Path(output_folder_path) / '2d_patches_training',
+#             stack.get_one_channel_data(patches_channel),
+#             zmask_meta,
+#             prefix=fstem,
+#             rescale_clip=0.001,
+#             make_3d=False,
+#             focus_metric='max_sobel',
+#         )
+#         ti.click('export_2d_patches')
+#
+#     if export_patch_masks and len(zmask_meta) > 0:
+#         files = export_patch_masks_from_zstack(
+#             Path(output_folder_path) / 'patch_masks',
+#             stack.get_one_channel_data(patches_channel),
+#             zmask_meta,
+#             prefix=fstem,
+#         )
+#
+#     if export_annotated_zstack:
+#         annotated = InMemoryDataAccessor(
+#             draw_boxes_on_3d_image(
+#                 stack.get_one_channel_data(patches_channel).data,
+#                 zmask_meta,
+#                 add_label=draw_label_on_zstack,
+#             )
+#         )
+#         write_accessor_data_to_file(
+#             Path(output_folder_path) / 'annotated_zstacks' / (fstem + '.tif'),
+#             annotated
+#         )
+#         ti.click('export_annotated_zstack')
+#
+#     # generate multichannel projection from label centroids
+#     dff = df[df['keeper']]
+#     if len(zmask_meta) > 0:
+#         interm['projected'] = project_stack_from_focal_points(
+#             dff['centroid-0'].to_numpy(),
+#             dff['centroid-1'].to_numpy(),
+#             dff['zi'].to_numpy(),
+#             stack,
+#             degree=4,
+#         )
+#     else: # else just return MIP
+#         interm['projected'] = stack.data.max(axis=-1)
+#
+#     return {
+#         'pixel_model_id': pixel_classifier.model_id,
+#         'input_filepath': input_file_path,
+#         'number_of_objects': len(zmask_meta),
+#         'pixeL_scale_in_micrometers': stack.pixel_scale_in_micrometers,
+#         'success': True,
+#         'timer_results': ti.events,
+#         'dataframe': df[df['keeper'] == True],
+#         'interm': interm,
+#     }
+
+def infer_object_map_from_zstack(
+        input_file_path: str,
+        output_folder_path: str,
+        models: List[Model],
+        pxmap_foreground_channel: int,
+        pxmap_threshold: float,
+        segmentation_channel: int,
+        patches_channel: int,
+        zmask_zindex: int = None,  # None for MIP,
+        zmask_clip: int = None,
+        zmask_type: str = 'boxes',
+        zmask_filters: Dict = None,
+        # zmask_expand_box_by: int = None,
+        exports: ZMaskExportParams = None,
+        **kwargs,
+) -> Dict:
+    assert len(models) == 2
+    pixel_classifier = models[0]
+    assert isinstance(pixel_classifier, IlastikPixelClassifierModel)
+    object_classifier = models[1]
+    assert isinstance(object_classifier, PatchStackObjectClassifier)
+
     ti = Timer()
     stack = generate_file_accessor(Path(input_file_path))
     fstem = Path(input_file_path).stem
@@ -49,16 +265,23 @@ def get_zmask_meta(
     mip = InMemoryDataAccessor(
         zmask_data,
     )
-    pxmap, _ = ilastik_pixel_classifier.infer(mip)
+    pxmap, _ = pixel_classifier.infer(mip)
     ti.click('infer_pixel_probability')
 
+    if exports.pixel_probabilities:
+        write_accessor_data_to_file(
+            Path(output_folder_path) / 'pixel_probabilities' / (fstem + '.tif'),
+            pxmap
+        )
+        ti.click('export_pixel_probability')
+
     obmask = InMemoryDataAccessor(
         pxmap.data > pxmap_threshold
     )
     ti.click('threshold_pixel_mask')
 
     # make zmask
-    zmask, zmask_meta, df, interm = build_zmask_from_object_mask(
+    obj_table = ZMaskObjectTable(
         obmask.get_one_channel_data(pxmap_foreground_channel),
         stack.get_one_channel_data(segmentation_channel),
         mask_type=zmask_type,
@@ -68,194 +291,26 @@ def get_zmask_meta(
     ti.click('generate_zmasks')
 
     # record pixel scale
-    df['pixel_scale_in_micrometers'] = float(stack.pixel_scale_in_micrometers.get('X'))
-
-    return ti, stack, fstem, obmask, pxmap, zmask, zmask_meta, df, interm
-
-
-# TODO: unpack and validate inputs
-def export_patches_from_multichannel_zstack(
-        input_file_path: str,
-        output_folder_path: str,
-        models: List[Model],
-        pxmap_threshold: float,
-        pxmap_foreground_channel: int,
-        segmentation_channel: int,
-        patches_channel: int,
-        zmask_zindex: int = None,  # None for MIP,
-        zmask_clip: int = None,
-        zmask_type: str = 'boxes',
-        zmask_filters: Dict = None,
-        zmask_expand_box_by: int = None,
-        export_pixel_probabilities=True,
-        export_2d_patches_for_training=True,
-        export_2d_patches_for_annotation=True,
-        draw_bounding_box_on_2d_patch=True,
-        draw_contour_on_2d_patch=False,
-        draw_mask_on_2d_patch=False,
-        export_3d_patches=True,
-        export_annotated_zstack=True,
-        draw_label_on_zstack=False,
-        export_patch_masks=True,
-        rgb_overlay_channels=(None, None, None),
-        rgb_overlay_weights=(1.0, 1.0, 1.0),
-) -> Dict:
-    pixel_classifier = models[0]
-
-    ti, stack, fstem, obmask, pxmap, zmask, zmask_meta, df, interm = get_zmask_meta(
-        input_file_path,
-        pixel_classifier,
-        segmentation_channel,
-        pxmap_threshold,
-        pxmap_foreground_channel=pxmap_foreground_channel,
-        zmask_zindex=zmask_zindex,
-        zmask_clip=zmask_clip,
-        zmask_expand_box_by=zmask_expand_box_by,
-        zmask_filters=zmask_filters,
-        zmask_type=zmask_type,
-    )
-
-    if export_pixel_probabilities:
-        write_accessor_data_to_file(
-            Path(output_folder_path) / 'pixel_probabilities' / (fstem + '.tif'),
-            pxmap
-        )
-        ti.click('export_pixel_probability')
-
-    if export_3d_patches and len(zmask_meta) > 0:
-        files = export_patches_from_zstack(
-            Path(output_folder_path) / '3d_patches',
-            stack.get_one_channel_data(patches_channel),
-            zmask_meta,
-            prefix=fstem,
-            draw_bounding_box=False,
-            rescale_clip=0.001,
-            make_3d=True,
-        )
-        ti.click('export_3d_patches')
-
-    if export_2d_patches_for_annotation and len(zmask_meta) > 0:
-        files = export_multichannel_patches_from_zstack(
-            Path(output_folder_path) / '2d_patches_annotation',
-            stack,
-            zmask_meta,
-            prefix=fstem,
-            rescale_clip=0.001,
-            make_3d=False,
-            focus_metric='max_sobel',
-            ch_white=patches_channel,
-            ch_rgb_overlay=rgb_overlay_channels,
-            draw_bounding_box=draw_bounding_box_on_2d_patch,
-            bounding_box_channel=1,
-            bounding_box_linewidth=2,
-            draw_contour=draw_contour_on_2d_patch,
-            draw_mask=draw_mask_on_2d_patch,
-            overlay_gain=rgb_overlay_weights,
-        )
-        df_patches = pd.DataFrame(files)
-        ti.click('export_2d_patches')
-        # associate 2d patches, dropping labeled objects that were not exported as patches
-        df = pd.merge(df, df_patches, left_index=True, right_on='df_index').drop(columns='df_index')
-        # prepopulate patch UUID
-        df['patch_id'] = df.apply(lambda _: uuid4(), axis=1)
-
-    if export_2d_patches_for_training and len(zmask_meta) > 0:
-        files = export_multichannel_patches_from_zstack(
-            Path(output_folder_path) / '2d_patches_training',
-            stack.get_one_channel_data(patches_channel),
-            zmask_meta,
-            prefix=fstem,
-            rescale_clip=0.001,
-            make_3d=False,
-            focus_metric='max_sobel',
-        )
-        ti.click('export_2d_patches')
-
-    if export_patch_masks and len(zmask_meta) > 0:
-        files = export_patch_masks_from_zstack(
-            Path(output_folder_path) / 'patch_masks',
-            stack.get_one_channel_data(patches_channel),
-            zmask_meta,
-            prefix=fstem,
-        )
-
-    if export_annotated_zstack:
-        annotated = InMemoryDataAccessor(
-            draw_boxes_on_3d_image(
-                stack.get_one_channel_data(patches_channel).data,
-                zmask_meta,
-                add_label=draw_label_on_zstack,
-            )
-        )
-        write_accessor_data_to_file(
-            Path(output_folder_path) / 'annotated_zstacks' / (fstem + '.tif'),
-            annotated
-        )
-        ti.click('export_annotated_zstack')
-
-    # generate multichannel projection from label centroids
-    dff = df[df['keeper']]
-    if len(zmask_meta) > 0:
-        interm['projected'] = project_stack_from_focal_points(
-            dff['centroid-0'].to_numpy(),
-            dff['centroid-1'].to_numpy(),
-            dff['zi'].to_numpy(),
-            stack,
-            degree=4,
-        )
-    else: # else just return MIP
-        interm['projected'] = stack.data.max(axis=-1)
-
-    return {
-        'pixel_model_id': pixel_classifier.model_id,
-        'input_filepath': input_file_path,
-        'number_of_objects': len(zmask_meta),
-        'pixeL_scale_in_micrometers': stack.pixel_scale_in_micrometers,
-        'success': True,
-        'timer_results': ti.events,
-        'dataframe': df[df['keeper'] == True],
-        'interm': interm,
-    }
-
-def infer_object_map_from_zstack(
-        input_file_path: str,
-        output_folder_path: str,
-        models: List[Model],
-        pxmap_threshold: float,
-        pxmap_foreground_channel: int,
-        segmentation_channel: int,
-        patches_channel: int,
-        zmask_zindex: int = None,  # None for MIP,
-        zmask_clip: int = None,
-        zmask_type: str = 'boxes',
-        zmask_filters: Dict = None,
-        # zmask_expand_box_by: int = None,
-        **kwargs,
-) -> Dict:
-    assert len(models) == 2
-    pixel_classifier = models[0]
-    assert isinstance(pixel_classifier, IlastikPixelClassifierModel)
-    object_classifier = models[1]
-    assert isinstance(object_classifier, PatchStackObjectClassifier)
-
-    ti, stack, fstem, obmask, pxmap, zmask, zmask_meta, df, interm = get_zmask_meta(
-        input_file_path,
-        pixel_classifier,
-        segmentation_channel,
-        pxmap_threshold,
-        pxmap_foreground_channel=pxmap_foreground_channel,
-        zmask_zindex=zmask_zindex,
-        zmask_clip=zmask_clip,
-        # zmask_expand_box_by=zmask_expand_box_by,
-        zmask_filters=zmask_filters,
-        zmask_type=zmask_type,
-        **kwargs
-    )
+    obj_table.df['pixel_scale_in_micrometers'] = float(stack.pixel_scale_in_micrometers.get('X'))
+
+    # ti, stack, fstem, obmask, pxmap, obj_table = get_zmask_meta(
+    #     input_file_path,
+    #     pixel_classifier,
+    #     segmentation_channel,
+    #     pxmap_threshold,
+    #     pxmap_foreground_channel=pxmap_foreground_channel,
+    #     zmask_zindex=zmask_zindex,
+    #     zmask_clip=zmask_clip,
+    #     # zmask_expand_box_by=zmask_expand_box_by,
+    #     zmask_filters=zmask_filters,
+    #     zmask_type=zmask_type,
+    #     **kwargs
+    # )
 
     # extract patches to accessor
     patches_acc = get_patches_from_zmask_meta(
         stack.get_one_channel_data(patches_channel),
-        zmask_meta,
+        obj_table.zmask_meta,
         rescale_clip=zmask_clip,
         make_3d=False,
         focus_metric='max_sobel',
@@ -265,22 +320,22 @@ def infer_object_map_from_zstack(
     # extract masks
     patch_masks_acc = get_patch_masks_from_zmask_meta(
         stack,
-        zmask_meta,
+        obj_table.zmask_meta,
         **kwargs
     )
 
     # send patches and mask stacks to object classifier
     result_acc, _ = object_classifier.infer(patches_acc, patch_masks_acc)
 
-    labels_map = interm['label_map']
+    labels_map = obj_table.interm['label_map']
     output_map = np.zeros(labels_map.shape, dtype=labels_map.dtype)
-    assert labels_map.shape == interm['label_map'].shape
-    assert labels_map.dtype == interm['label_map'].dtype
+    assert labels_map.shape == obj_table.get_label_map().shape
+    assert labels_map.dtype == obj_table.get_label_map().dtype
 
     # assign labels to object map:
     meta = []
-    for ii in range(0, len(zmask_meta)):
-        object_id = zmask_meta[ii]['info'].label
+    for ii in range(0, len(obj_table.zmask_meta)):
+        object_id = obj_table.zmask_meta[ii]['info'].label
         result_patch = mask_largest_object(result_acc.iat(ii))
         object_class = np.unique(result_patch)[1]
         output_map[labels_map == object_id] = object_class
@@ -293,6 +348,51 @@ def infer_object_map_from_zstack(
     )
     ti.click('export_object_classes')
 
+    if exports.patches_3d:
+        obj_table.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:
+        obj_table.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:
+        obj_table.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:
+        obj_table.export_patch_masks(
+            Path(output_folder_path) / 'patch_masks',
+            fstem,
+            patches_channel,
+            exports.patch_masks
+        )
+
+    if exports.annotated_z_stack:
+        obj_table.export_annotated_zstack(
+            Path(output_folder_path) / 'patch_masks',
+            fstem,
+            patches_channel,
+            exports.annotated_z_stack
+        )
+        ti.click('export_annotated_zstack')
+
+
     return {
         'timer_results': ti.events,
         'dataframe': pd.DataFrame(meta),

extensions/chaeo/zmask.py

+from uuid import uuid4
+
 import numpy as np
 import pandas as pd
 
@@ -5,7 +7,149 @@ from skimage.measure import find_contours, label, regionprops_table
 from sklearn.preprocessing import PolynomialFeatures
 from sklearn.linear_model import LinearRegression
 
-from model_server.accessors import GenericImageDataAccessor
+from extensions.chaeo.annotators import draw_boxes_on_3d_image
+from extensions.chaeo.products import export_patches_from_zstack, export_multichannel_patches_from_zstack, export_patch_masks_from_zstack
+from extensions.chaeo.params import ZMaskExportParams
+from model_server.accessors import GenericImageDataAccessor, InMemoryDataAccessor, write_accessor_data_to_file
+
+
+
+class ZMaskObjectTable(object):
+
+    def __init__(
+            self,
+            acc_mask: GenericImageDataAccessor,
+            acc_raw: GenericImageDataAccessor,
+            filters=None,
+            mask_type='contours',
+            expand_box_by=(0, 0),
+    ):
+        self.zmask, self.zmask_meta, self.df, self.interm = build_zmask_from_object_mask(
+            acc_mask,
+            acc_raw,
+            filters=filters,
+            mask_type=mask_type,
+            expand_box_by=expand_box_by
+        )
+        self.acc_raw = acc_raw
+        self.count = len(self.zmask_meta)
+
+    def get_label_map(self):
+        return self.interm.lamap
+
+    def get_argmax(self):
+        return self.interm.argmax
+
+    def export_3d_patches(self, where, prefix, channel, params: ZMaskExportParams):
+        if not self.count:
+            return
+        files = export_patches_from_zstack(
+            where,
+            self.acc_raw.get_one_channel_data(channel),
+            self.zmask_meta,
+            prefix=prefix,
+            draw_bounding_box=params.draw_bounding_box,
+            rescale_clip=params.rescale_clip,
+            make_3d=True,
+        )
+
+    def export_2d_patches_for_annotation(self, where, prefix, channel, params: ZMaskExportParams):
+        if not self.count:
+            return
+        files = export_multichannel_patches_from_zstack(
+            where,
+            self.acc_raw.get_one_channel_data(channel),
+            self.zmask_meta,
+            prefix=prefix,
+            draw_bounding_box=params.draw_bounding_box,
+            rescale_clip=params.rescale_clip,
+            make_3d=False,
+            focus_metric=params.focus_metric,
+            ch_white=channel,
+            ch_rgb_overlay=params.rgb_overlay_channels,
+            bounding_box_channel=1,
+            bounding_box_linewidth=2,
+            draw_contour=params.draw_contour,
+            draw_mask=params.draw_mask,
+            overlay_gain=params.rgb_overlay_weights,
+        )
+        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_training(self, where, prefix, channel, params: ZMaskExportParams):
+        if not self.count:
+            return
+        files = export_multichannel_patches_from_zstack(
+            where,
+            self.acc_raw.get_one_channel_data(channel),
+            self.zmask_meta,
+            prefix=prefix,
+            draw_bounding_box=params.draw_bounding_box,
+            rescale_clip=params.rescale_clip,
+            make_3d=False,
+            focus_metric=params.focus_metric,
+            ch_white=channel,
+            ch_rgb_overlay=params.rgb_overlay_channels,
+            bounding_box_channel=1,
+            bounding_box_linewidth=2,
+            draw_contour=params.draw_contour,
+            draw_mask=params.draw_mask,
+            overlay_gain=params.rgb_overlay_weights,
+        )
+        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: ZMaskExportParams):
+        if not self.count:
+            return
+        files = export_multichannel_patches_from_zstack(
+            where,
+            self.acc_raw.get_one_channel_data(channel),
+            self.zmask_meta,
+            prefix=prefix,
+            rescale_clip=params.rescale_clip,
+            make_3d=False,
+            focus_metric=params.focus_metric,
+        )
+
+    def export_patch_masks(self, where, prefix, channel, params: ZMaskExportParams):
+        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: ZMaskExportParams):
+        annotated = InMemoryDataAccessor(
+            draw_boxes_on_3d_image(
+                self.acc_raw.get_one_channel_data(channel).data,
+                self.zmask_meta,
+                add_label=params.draw_label,
+            )
+        )
+        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:
+            projected = project_stack_from_focal_points(
+                dff['centroid-0'].to_numpy(),
+                dff['centroid-1'].to_numpy(),
+                dff['zi'].to_numpy(),
+                self.acc_raw,
+                degree=4,
+            )
+        else:  # else just return MIP
+            projected = self.acc_raw.data.max(axis=-1)
+        return projected
 
 def build_zmask_from_object_mask(
         obmask: GenericImageDataAccessor,

extensions/ilastik/models.py

@@ -53,6 +53,7 @@ class IlastikImageToImageModel(ImageToImageModel):
 
 class IlastikPixelClassifierModel(IlastikImageToImageModel):
     model_id = 'ilastik_pixel_classification'
+    operations = ['segment', ]
 
     @staticmethod
     def get_workflow():
@@ -77,35 +78,40 @@ class IlastikPixelClassifierModel(IlastikImageToImageModel):
         )
         return InMemoryDataAccessor(data=yxcz), {'success': True}
 
-class IlastikObjectClassifierFromPixelPredictionsModel(IlastikImageToImageModel):
-    model_id = 'ilastik_object_classification_from_pixel_predictions'
-
-    @staticmethod
-    def get_workflow():
-        from ilastik.workflows.objectClassification.objectClassificationWorkflow import ObjectClassificationWorkflowPrediction
-        return ObjectClassificationWorkflowPrediction
-
-    def infer(self, input_img: GenericImageDataAccessor, pxmap_img: GenericImageDataAccessor) -> (np.ndarray, dict):
-        tagged_input_data = vigra.taggedView(input_img.data, 'yxcz')
-        tagged_pxmap_data = vigra.taggedView(pxmap_img.data, 'yxcz')
-
-        dsi = [
-            {
-                'Raw Data': self.PreloadedArrayDatasetInfo(preloaded_array=tagged_input_data),
-                'Prediction Maps': self.PreloadedArrayDatasetInfo(preloaded_array=tagged_pxmap_data),
-            }
-        ]
-
-        obmaps = self.shell.workflow.batchProcessingApplet.run_export(dsi, export_to_array=True) # [z x h x w x n]
-
-        assert (len(obmaps) == 1, 'ilastik generated more than one object map')
-
-        yxcz = np.moveaxis(
-            obmaps[0],
-            [1, 2, 3, 0],
-            [0, 1, 2, 3]
+    def segment(self, input_img, thresh, channel):
+        return InMemoryDataAccessor(
+            self.infer(input_img).data[:, :, channel, :] > thresh
         )
-        return InMemoryDataAccessor(data=yxcz), {'success': True}
+
+# class IlastikObjectClassifierFromPixelPredictionsModel(IlastikImageToImageModel):
+#     model_id = 'ilastik_object_classification_from_pixel_predictions'
+#
+#     @staticmethod
+#     def get_workflow():
+#         from ilastik.workflows.objectClassification.objectClassificationWorkflow import ObjectClassificationWorkflowPrediction
+#         return ObjectClassificationWorkflowPrediction
+#
+#     def infer(self, input_img: GenericImageDataAccessor, pxmap_img: GenericImageDataAccessor) -> (np.ndarray, dict):
+#         tagged_input_data = vigra.taggedView(input_img.data, 'yxcz')
+#         tagged_pxmap_data = vigra.taggedView(pxmap_img.data, 'yxcz')
+#
+#         dsi = [
+#             {
+#                 'Raw Data': self.PreloadedArrayDatasetInfo(preloaded_array=tagged_input_data),
+#                 'Prediction Maps': self.PreloadedArrayDatasetInfo(preloaded_array=tagged_pxmap_data),
+#             }
+#         ]
+#
+#         obmaps = self.shell.workflow.batchProcessingApplet.run_export(dsi, export_to_array=True) # [z x h x w x n]
+#
+#         assert (len(obmaps) == 1, 'ilastik generated more than one object map')
+#
+#         yxcz = np.moveaxis(
+#             obmaps[0],
+#             [1, 2, 3, 0],
+#             [0, 1, 2, 3]
+#         )
+#         return InMemoryDataAccessor(data=yxcz), {'success': True}
 
 
 class IlastikObjectClassifierFromSegmentationModel(IlastikImageToImageModel):

extensions/ilastik/router.py

@@ -35,7 +35,11 @@ def load_ilastik_model(model_class: ilm.IlastikImageToImageModel, project_file:
         )
     return {'model_id': result}
 
-@router.put('/px/load/')
+# @router.put('/px/load/')
+# def load_px_model(project_file: str, duplicate: bool = True) -> dict:
+#     return load_ilastik_model(ilm.IlastikPixelClassifierModel, project_file, duplicate=duplicate)
+
+@router.put('/seg/load/')
 def load_px_model(project_file: str, duplicate: bool = True) -> dict:
     return load_ilastik_model(ilm.IlastikPixelClassifierModel, project_file, duplicate=duplicate)