workflows.py

"""
Implementation of image analysis work behind API endpoints, without knowledge of persistent data in server session.
"""
from pathlib import Path
from typing import Dict

from extensions.ilastik.models import IlastikPixelClassifierModel, IlastikObjectClassifierFromPixelPredictionsModel
from model_server.accessors import generate_file_accessor, write_accessor_data_to_file
from model_server.workflows import Timer

from pydantic import BaseModel

class WorkflowRunRecord(BaseModel):
    pixel_model_id: str
    object_model_id: str
    input_filepath: str
    pixel_map_filepath: str
    object_map_filepath: str
    success: bool
    timer_results: Dict[str, float]


def infer_px_then_ob_model(
        fpi: Path,
        px_model: IlastikPixelClassifierModel,
        ob_model: IlastikObjectClassifierFromPixelPredictionsModel,
        where_output: Path,
        **kwargs
) -> WorkflowRunRecord:
    """
    Workflow that specifically runs an ilastik pixel classifier, then passes results to an object classifier,
    saving intermediate images
    :param fpi: Path object that references input image file
    :param px_model: model instance for pixel classification
    :param ob_model: model instance for object classification
    :param where_output: Path object that references output image directory
    :param kwargs: variable-length keyword arguments
    :return:
    """
    assert isinstance(px_model, IlastikPixelClassifierModel)
    assert isinstance(ob_model, IlastikObjectClassifierFromPixelPredictionsModel)

    ti = Timer()
    ch = kwargs.get('channel')
    img = generate_file_accessor(fpi).get_one_channel_data(ch)
    ti.click('file_input')

    px_map, _ = px_model.infer(img)
    ti.click('pixel_probability_inference')

    px_map_path = where_output / (px_model.model_id + '_pxmap_' + fpi.stem + '.tif')
    write_accessor_data_to_file(px_map_path, px_map)
    ti.click('pixel_map_output')

    ob_map, _ = ob_model.infer(img, px_map)
    ti.click('object_classification')

    ob_map_path = where_output / (ob_model.model_id + '_obmap_' + fpi.stem + '.tif')
    write_accessor_data_to_file(ob_map_path, ob_map)
    ti.click('object_map_output')

    return WorkflowRunRecord(
        pixel_model_id=px_model.model_id,
        object_model_id=ob_model.model_id,
        input_filepath=str(fpi),
        pixel_map_filepath=str(px_map_path),
        object_map_filepath=str(ob_map_path),
        success=True,
        timer_results=ti.events,
    )