workflow.py

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

from model_server.image import generate_file_accessor, write_accessor_data_to_file
from model_server.model import Model

from pydantic import BaseModel

class Timer(object):
    tfunc = perf_counter

    def __init__(self):
        self.events = {}
        self.last = self.tfunc()

    def click(self, key):
        self.events[key] = self.tfunc() - self.last
        self.last = self.tfunc()

class WorkflowRunRecord(BaseModel):
    model_id: str
    input_filepath: str
    output_filepath: str
    success: bool
    timer_results: Dict[str, float]


def infer_image_to_image(fpi: Path, model: Model, where_output: Path, **kwargs) -> WorkflowRunRecord:
    """
    Generic workflow where a model processes an input image into an output image
    :param fpi: Path object that references input image file
    :param model: model object
    :param where_output: Path object that references output image directory
    :param kwargs: variable-length keyword arguments
    :return: record object
    """
    ti = Timer()
    ch = kwargs.get('channel')
    img = generate_file_accessor(fpi).get_one_channel_data(ch)
    ti.click('file_input')

    outdata, _ = model.infer(img)
    ti.click('inference')

    outpath = where_output / (model.model_id + '_' + fpi.stem + '.tif')
    write_accessor_data_to_file(outpath, outdata)
    ti.click('file_output')

    return WorkflowRunRecord(
        model_id=model.model_id,
        input_filepath=str(fpi),
        output_filepath=str(outpath),
        success=True,
        timer_results=ti.events,
    )