Prototyped object classifier with segmentation

extensions/ilastik/models.py

@@ -106,3 +106,35 @@ class IlastikObjectClassifierFromPixelPredictionsModel(IlastikImageToImageModel)
             [0, 1, 2, 3]
         )
         return InMemoryDataAccessor(data=yxcz), {'success': True}
+
+
+class IlastikObjectClassifierFromSegmentationModel(IlastikImageToImageModel):
+    model_id = 'ilastik_object_classification_from_segmentation'
+
+    @staticmethod
+    def get_workflow():
+        from ilastik.workflows.objectClassification.objectClassificationWorkflow import ObjectClassificationWorkflowBinary
+        return ObjectClassificationWorkflowBinary
+
+    def infer(self, input_img: GenericImageDataAccessor, segmentation_img: GenericImageDataAccessor) -> (np.ndarray, dict):
+        assert segmentation_img.is_mask()
+        tagged_input_data = vigra.taggedView(input_img.data, 'yxcz')
+        tagged_seg_data = vigra.taggedView(segmentation_img.data, 'yxcz')
+
+        dsi = [
+            {
+                'Raw Data': self.PreloadedArrayDatasetInfo(preloaded_array=tagged_input_data),
+                'Segmentation Image': self.PreloadedArrayDatasetInfo(preloaded_array=tagged_seg_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}

extensions/ilastik/tests/test_ilastik.py

@@ -199,3 +199,5 @@ class TestIlastikOverApi(TestServerBaseClass):
             }
         )
         self.assertEqual(resp_infer.status_code, 200, resp_infer.content.decode())
+
+    # TODO: test IlastikObjectClassifierFromSegmentationModel when a test model is complete
\ No newline at end of file