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/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)