process all the images in an OME-NGFF plate, segmenting them and adding OME-NGFF labels to each image

ngff_plate_segment_add_labels.py

import os
import zarr

from skimage.filters import threshold_otsu
from skimage.segmentation import clear_border
from skimage.measure import label
from skimage.morphology import closing, square

from ome_zarr.io import parse_url
from ome_zarr.writer import write_labels

plate_path = "path/to/plate.zarr"
location = parse_url(plate_path, mode="r")
print('location.root_attrs', location.root_attrs)

wells_json = location.root_attrs["plate"]["wells"]

# Iterate through each Well in Plate...
for well_json in wells_json:
    path = well_json["path"]

    well_path = os.path.join(plate_path, path)
    print("well_path", well_path)
    well = parse_url(well_path, mode="r")

    for image_json in well.root_attrs["well"]["images"]:
        image_path = os.path.join(well_path, image_json["path"])
        print("image_path", image_path)

        zarr_img = parse_url(image_path)
        multiscale = zarr_img.root_attrs["multiscales"][0]
        img_axes = multiscale["axes"]

        # read full-size resolution
        image_3d = zarr_img.load(multiscale["datasets"][0]["path"])
        print("image_3d", image_3d.shape)

        # FIXME - hard-coded for CYX images
        channel_index = 0
        image = image_3d[channel_index]
        label_axes = img_axes[-2:]

        # dask.compute() -> numpy array
        image = image.compute()

        # https://scikit-image.org/docs/stable/auto_examples/segmentation/plot_label.html#sphx-glr-auto-examples-segmentation-plot-label-py
        # apply threshold
        thresh = threshold_otsu(image)
        bw = closing(image > thresh, square(3))
        # remove artifacts connected to image border
        cleared = clear_border(bw)
        # label image regions
        label_image = label(cleared)

        # write labels under image/labels/ ...
        img_store = parse_url(image_path, mode="w").store
        img_group = zarr.group(store=img_store)
        label_name = "segmentation"
        write_labels(
            label_image,
            img_group,
            name=label_name,
            axes=label_axes
        )