Mocked up ilastik model classes

model_server/ilastik.py

+import os
+
+from ilastik import app
+from ilastik.applets.dataSelection import DatasetInfo, FilesystemDatasetInfo
+from ilastik.applets.dataSelection.opDataSelection import PreloadedArrayDatasetInfo
+from ilastik.workflows.pixelClassification import PixelClassificationWorkflow
+from ilastik.workflows.objectClassification.objectClassificationWorkflow import ObjectClassificationWorkflow
+
+import numpy as np
+
+from model_server.model import ImageToImageModel, ParameterExpectedError
+
+
+class IlastikImageToImageModel(ImageToImageModel):
+
+    # workflows = {
+    #     'pixel_classification': PixelClassificationWorkflow,
+    #     'object_classification': ObjectClassificationWorkflow,
+    # }
+
+    def __init__(self, params, autoload=True):
+        super().__init__(autoload, params)
+
+        if 'project_file' not in params or not os.path.exists(params['project_file']):
+            raise ParameterExpectedError('Ilastik model expects a project (*.ilp) file')
+        self.project_file = params['project_file']
+
+    def load(self):
+        os.environ["LAZYFLOW_THREADS"] = "8"
+        os.environ["LAZYFLOW_TOTAL_RAM_MB"] = "24000"
+
+        args = app.parse_args([])
+        args.headless = True
+        args.project = self.project_file
+        shell = app.main(args)
+
+        if shell.workflow.__class__ != 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.shell = shell
+
+    @staticmethod
+    def get_top_level_operator(workflow):
+        if isinstance(workflow, PixelClassificationWorkflow):
+            return workflow.pcApplet.topLevelOperator
+        elif isinstance(workflow, ObjectClassificationWorkflow):
+            return workflow.objectClassificationApplet.topLevelOperator
+
+
+    # def infer(self, img, channel=None) -> (np.ndarray, dict):
+    #     assert self.operator.Classifier.ready()
+
+class IlastikPixelClassifierModel(IlastikImageToImageModel):
+    workflow = PixelClassificationWorkflow
+
+    def infer(self, input_img: np.ndarray, channel=None) -> (np.ndarray, dict):
+        dsi = [
+            {
+                'Raw Data': PreloadedArrayDatasetInfo(preloaded_array=input_img),
+            }
+        ]
+        pxmap = self.shell.workflow.batchProcessingApplet.run_export(dsi, export_to_array=True)
+        return pxmap
+
+
+