workflow.py

"""
Implementation of image analysis work behind API endpoints, without knowledge of persistent data in server session.
"""

from time import time
from typing import Dict

from model_server.image import generate_file_accessor, WriteableTiffFileAccessor

from pydantic import BaseModel

# TODO: timer decorator

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

def infer_image_to_image(fpi, model, where_output, **kwargs) -> dict:
    # read image file into memory
    # maybe this isn't accurate if e.g. czifile loads lazily
    t0 = time()
    img = generate_file_accessor(fpi)
    dt_fi = time() - t0

    # TODO: assert indata format
    assert(img.shape_dict['T'] == 1)
    assert (img.shape_dict['T'] == 1)

    # run model inference
    # TODO: call this async / await and report out infer status to optional callback
    ch = kwargs.get('channel')
    outdata, messages = model.infer(img, channel=ch)
    dt_inf = time() - t0

    # TODO: assert outdata format

    # write output to file
    outpath = where_output / (img.fpath.stem + '.tif')
    WriteableTiffFileAccessor(outpath).write(outdata)
    dt_fo = time() - t0

    # TODO: smoother step-timing e.g. w/ decorate
    timer_results = {
        'file_input': dt_fi,
        'inference': dt_inf - dt_fi,
        'file_output': dt_fo - dt_fi - dt_inf
    }

    record = WorkflowRunRecord(
        model_id=model.model_id,
        input_filepath=str(fpi),
        output_filepath=str(outpath),
        success=messages['success'],
        timer_results=timer_results
    )

    return record