Refined workflow class check

conf/testing.py

@@ -13,6 +13,7 @@ czifile = {
 
 ilastik = {
     'pixel_classifier': root / 'testdata' / 'pix02.ilp',
+    'object_classifier': root / 'testdata' / 'obj01.ilp',
 }
 
 output_path = root / 'testing_output'

model_server/ilastik.py

@@ -20,7 +20,7 @@ class IlastikImageToImageModel(ImageToImageModel):
             raise ParameterExpectedError('Ilastik model expects a project (*.ilp) file')
         self.project_file = str(params['project_file'])
         self.shell = None
-        self.operator = None
+        # self.operator = None
         super().__init__(autoload, params)
 
     def __del__(self):
@@ -35,27 +35,27 @@ class IlastikImageToImageModel(ImageToImageModel):
         args.project = self.project_file
         shell = app.main(args)
 
-        if shell.workflow.__class__ != self.workflow:
+        if not isinstance(shell.workflow, self.workflow):
             raise ParameterExpectedError(
                 f'Ilastik project file {self.project_file} does not describe an instance of {shell.workflow.__class__}'
             )
-        self.operator = self.get_top_level_operator(shell.workflow)
+        # self.operator = self.get_top_level_operator(shell.workflow)
         self.shell = shell
 
         return True
 
-    @staticmethod
-    def get_top_level_operator(workflow):
-        if isinstance(workflow, PixelClassificationWorkflow):
-            return workflow.pcApplet.topLevelOperator
-        elif isinstance(workflow, ObjectClassificationWorkflow):
-            return workflow.objectClassificationApplet.topLevelOperator
+    # @staticmethod
+    # def get_top_level_operator(workflow):
+    #     if isinstance(workflow, PixelClassificationWorkflow):
+    #         return workflow.pcApplet.topLevelOperator
+    #     elif isinstance(workflow, ObjectClassificationWorkflow):
+    #         return workflow.objectClassificationApplet.topLevelOperator
 
 
 class IlastikPixelClassifierModel(IlastikImageToImageModel):
     workflow = PixelClassificationWorkflow
 
-    def infer(self, input_img: GenericImageDataAccessor, channel=None) -> (np.ndarray, dict):
+    def infer(self, input_img: GenericImageDataAccessor) -> (np.ndarray, dict):
         tagged_input_data = vigra.taggedView(input_img.data, 'xycz')
         dsi = [
             {
@@ -64,7 +64,7 @@ class IlastikPixelClassifierModel(IlastikImageToImageModel):
         ]
         pxmaps = self.shell.workflow.batchProcessingApplet.run_export(dsi, export_to_array=True) # [1 x w x h x n]
 
-        assert(len(pxmaps) == 1, 'ilastik generated more than on pixel map')
+        assert(len(pxmaps) == 1, 'ilastik generated more than one pixel map')
 
         xycz = np.moveaxis(
             pxmaps[0],
@@ -73,5 +73,27 @@ class IlastikPixelClassifierModel(IlastikImageToImageModel):
         )
         return InMemoryDataAccessor(data=xycz), {'success': True}
 
+class IlastikObjectClassifierModel(IlastikImageToImageModel):
+    workflow = ObjectClassificationWorkflow
 
+    def infer(self, input_img: GenericImageDataAccessor, pxmap_img: GenericImageDataAccessor) -> (np.ndarray, dict):
+        tagged_input_data = vigra.taggedView(input_img.data, 'xycz')
+        tagged_pxmap_data = vigra.taggedView(pxmap_img.data, 'xycz')
+
+        dsi = [
+            {
+                'Raw Data': PreloadedArrayDatasetInfo(preloaded_array=tagged_input_data),
+                'Prediction Maps': PreloadedArrayDatasetInfo(preloaded_array=tagged_pxmap_data),
+            }
+        ]
 
+        obmaps = self.shell.workflow.batchProcessingApplet.run_export(dsi, export_to_array=True)
+
+        assert (len(obmaps) == 1, 'ilastik generated more than one object map')
+
+        xycz = np.moveaxis(
+            obmaps[0],
+            [2, 1, 3, 0],
+            [0, 1, 2, 3]
+        )
+        return InMemoryDataAccessor(data=xycz), {'success': True}

tests/test_ilastik.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from conf.testing import czifile, ilastik, output_path
 from model_server.image import CziImageFileAccessor, InMemoryDataAccessor, write_accessor_data_to_file
-from model_server.ilastik import IlastikPixelClassifierModel
+from model_server.ilastik import IlastikObjectClassifierModel, IlastikPixelClassifierModel
 
 class TestIlastikPixelClassification(unittest.TestCase):
     def setUp(self) -> None:
@@ -65,3 +65,6 @@ class TestIlastikPixelClassification(unittest.TestCase):
             )
         )
 
+    def test_run_object_classifier(self):
+        model = IlastikObjectClassifierModel({'project_file': ilastik['object_classifier']})
+