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
import vigra

from model_server.image import GenericImageDataAccessor, InMemoryDataAccessor
from model_server.model import ImageToImageModel, ParameterExpectedError


class IlastikImageToImageModel(ImageToImageModel):

    def __init__(self, params, autoload=True):
        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 = str(params['project_file'])
        self.shell = None
        self.operator = None
        super().__init__(autoload, params)

    def __del__(self):
        self.shell.closeCurrentProject()

    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

        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


class IlastikPixelClassifierModel(IlastikImageToImageModel):
    workflow = PixelClassificationWorkflow

    def infer(self, input_img: GenericImageDataAccessor, channel=None) -> (np.ndarray, dict):
        tagged_input_data = vigra.taggedView(input_img.data, 'xycz')
        dsi = [
            {
                'Raw Data': PreloadedArrayDatasetInfo(preloaded_array=tagged_input_data),
            }
        ]
        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')

        xycz = np.moveaxis(
            pxmaps[0],
            [2, 1, 3, 0],
            [0, 1, 2, 3]
        )
        return InMemoryDataAccessor(data=xycz), {'success': True}