zmask.py

import numpy as np
import pandas as pd

from skimage.measure import find_contours, label, regionprops_table

from model_server.accessors import GenericImageDataAccessor

def build_zmask_from_object_mask(
        obmask: GenericImageDataAccessor,
        zstack: GenericImageDataAccessor,
        filters=None,
        mask_type='contour',
        expand_box_by=(0, 0),
):
    """
    Given a 2D mask of objects, build a 3D mask, where each object's z-position is determined by the index of
    maximum intensity in z.  Return this zmask and a list of each object's meta information.
    :param obmask: GenericImageDataAccessor monochrome 2D inary mask of objects
    :param zstack: GenericImageDataAccessor monochrome zstack of same Y, X dimension as obmask
    :param filters: dictionary of form {attribute: (min, max)}; valid attributes are 'area' and 'solidity'
    :param mask_type: if 'boxes', zmask is True in each object's complete bounding box; otherwise 'contours'
    :param expand_box_by: (xy, z) expands bounding box by (xy, z) pixels except where this hits a boundary
    :return: tuple (zmask, meta)
        np.ndarray zmask: boolean mask of same size as stack
        meta: List containing one Dict per object, with keys:
            info: object's properties from skimage.measure.regionprops_table, including bounding box (y0, y1, x0, x1)
            slice: named slice (np.s_) of (optionally) expanded bounding box
            relative_bounding_box: bounding box (y0, y1, x0, x1) in relative frame of (optionally) expanded bounding box
            contour: object's contour returned by skimage.measure.find_contours
            mask: mask of object in relative frame of (optionally) expanded bounding box
    """

    # validate inputs
    assert zstack.chroma == 1
    assert zstack.nz > 1
    assert mask_type in ('contours', 'boxes'), mask_type
    assert obmask.is_mask()
    assert obmask.chroma == 1
    assert obmask.nz == 1
    assert zstack.hw == obmask.hw

    # assign object labels and build object query
    lamap = label(obmask.data[:, :, 0, 0])
    query_str = 'label > 0'  # always true
    if filters is not None:
        for k in filters.keys():
            assert k in ('area', 'solidity')
            vmin, vmax = filters[k]
            assert vmin >= 0
            query_str = query_str + f' & {k} > {vmin} & {k} < {vmax}'

    # build dataframe of objects, assign z index to each object
    argmax = zstack.data.argmax(axis=3, keepdims=True)[:, :, 0, 0]
    df = (
        pd.DataFrame(
            regionprops_table(
                lamap,
                intensity_image=argmax,
                properties=('label', 'area', 'intensity_mean', 'solidity', 'bbox')
            )
        )
        .query(query_str)
        .rename(
            columns={
                'bbox-0': 'y0',
                'bbox-1': 'x0',
                'bbox-2': 'y1',
                'bbox-3': 'x1',
            }
        )
    )
    df['zi'] = df['intensity_mean'].round().astype('int')

    # make an object map where label is replaced by focus position in stack and background is -1
    lut = np.zeros(lamap.max() + 1) - 1
    lut[df.label] = df.zi

    # convert bounding boxes to numpy slice objects
    ebxy, ebz = expand_box_by
    h, w, c, nz = zstack.shape

    meta = []
    for ob in df.itertuples(name='LabeledObject'):
        y0 = max(ob.y0 - ebxy, 0)
        y1 = min(ob.y1 + ebxy, h - 1)
        x0 = max(ob.x0 - ebxy, 0)
        x1 = min(ob.x1 + ebxy, w - 1)
        z0 = max(ob.zi - ebz, 0)
        z1 = min(ob.zi + ebz, nz)

        # relative bounding box positions
        rbb = {
            'y0': ob.y0 - y0,
            'y1': ob.y1 - y0,
            'x0': ob.x0 - x0,
            'x1': ob.x1 - x0,
        }

        sl = np.s_[y0: y1, x0: x1, :, z0: z1 + 1]

        # compute contours
        obmask = (lamap == ob.label)
        contour = find_contours(obmask)
        mask = obmask[ob.y0: ob.y1, ob.x0: ob.x1]

        meta.append({
            'info': ob,
            'slice': sl,
            'relative_bounding_box': rbb,
            'contour': contour,
            'mask': mask
        })

    # build mask z-stack
    zi_st = np.zeros(zstack.shape, dtype='bool')
    if mask_type == 'contours':
        zi_map = (lut[lamap] + 1.0).astype('int')
        idxs = np.array(zi_map) - 1
        np.put_along_axis(
            zi_st,
            np.expand_dims(idxs, (2, 3)),
            1,
            axis=3
        )

        # change background level from to 0 in final frame
        zi_st[:, :, :, -1][lamap == 0] = 0

    elif mask_type == 'boxes':
        for bb in meta:
            sl = bb['slice']
            zi_st[sl] = 1

    return zi_st, meta