convert_to_ngff_segmentation_v2.py

import os, dataclasses, numcodecs, abc, time, dask
from aicsimageio import AICSImage
from aicsimageio.metadata.utils import OME
import numpy as np, cupy as cp
from pathlib import Path
import zarr
from typing import (Union, Iterable)
import warnings
from dask import array as da, bag, delayed
from dask.highlevelgraph import HighLevelGraph
import dask
from dask_image import ndfilters
from dask_cuda import LocalCUDACluster
#from rmm.allocators.cupy import rmm_cupy_allocator
#import rmm

import itertools
from pathlib import Path
import glob, zarr
from multiscales import Multimeta
import utils_zarrp
from utils_zarrp import otsu

from typing import Callable, Any
from collections import defaultdict

from cellpose import core, models, metrics

from distributed import LocalCluster, Client
from joblib import delayed as jdel, Parallel, parallel_config
from joblib.externals.loky import get_reusable_executor
get_reusable_executor().shutdown()


def get_regions(array_shape,
                region_shape,
                as_slices = False
                ):
    assert len(array_shape) == len(region_shape)
    steps = []
    for i in range(len(region_shape)):
        size = array_shape[i]
        inc = region_shape[i]
        seq = np.arange(0, size, inc)
        if size > seq[-1]:
            seq = np.append(seq, size)
        increments = tuple([(seq[i], seq[i+1]) for i in range(len(seq) - 1)])
        tuples = tuple(tuple(item) for item in increments)
        if as_slices:
            slcs = tuple([slice(*item) for item in tuples])
            steps.append(slcs)
        else:
            steps.append(tuples)
    out = list(itertools.product(*steps))
    return out

def read_image(file_path: str):
    img = AICSImage(file_path)
    return img.get_image_dask_data()

def create_zarr_array(directory: str,
                   array_name: str,
                   shape: tuple,
                   chunks: tuple,
                   dtype: Any,
                   overwrite: bool = False,
                   ) -> zarr.Array:
    chunks = tuple(np.minimum(shape, chunks))
    if not isinstance(directory, zarr.Group):
        path = os.path.join(directory, array_name)
        dataset = zarr.create(shape=shape,
                              chunks=chunks,
                              dtype=dtype,
                              store=path,
                              dimension_separator='/',
                              overwrite=overwrite
                              )
    else:
        _ = directory.create(name = array_name,
                              shape = shape,
                              chunks = chunks,
                              dtype = dtype,
                              dimension_separator='/',
                              overwrite=overwrite
                              )
        dataset = directory[array_name]
    return dataset

def write_single_region(region: da.Array,
                dataset: str,
                region_slice: slice = None
                ):
    da.to_zarr(region,
               url = dataset,
               region = region_slice,
               compute=True,
               return_stored=True
               )
    return dataset

def write_regions_sequential(
                  image_regions: tuple,
                  region_slices: tuple,
                  dataset: zarr.Array,
                  **kwargs
                  ):
    executor = get_reusable_executor(max_workers=n_jobs,
                                     kill_workers=True,
                                     context='loky')
    for region_slice, image_region in zip(region_slices, image_regions):
        executor.submit(write_single_region,
                         region=image_region,
                         dataset=dataset,
                         region_slice=region_slice
                         )
    return dataset

def write_regions(
                    image_regions: tuple,
                    region_slices: tuple,
                    dataset: zarr.Array,
                    client: Client = None
                    ) -> zarr.Array:
    if client is None:
        n_jobs = 4
    else:
        n_jobs = client.cluster.workers.__len__()
    client.cluster.scale(n_jobs)
    client.scatter(image_regions)
    client.scatter(region_slices)
    with parallel_config(backend = 'loky', n_jobs = n_jobs):
        with Parallel() as parallel:
            parallel(jdel(write_single_region)(region = image_region,
                                               region_slice = region_slice,
                                               dataset = dataset)
                     for image_region, region_slice in
                     zip(image_regions, region_slices)
                     )
    return dataset


def normalize_custom(img,
    lower=0.0,
    upper=400.0,
    background=100,
    percentile = False,
    percentile_lower = 1,
    percentile_upper = 99,
    ):

    if percentile:
        x01 = np.percentile(img, percentile_lower)
        x99 = np.percentile(img, percentile_upper)
        if x99 - x01 > 1e-3:
            img = (img - x01) / (x99 - x01)
        else:
            img[:] = 0
    else:
        img = np.asarray(img)
        img = img.astype(np.float32)
        img = img - background
        img[img<0] = 0
        img = (img - lower) / (upper - lower)

    return img

def process_and_save_to_ngff(arr: da.Array,
                            output_path: str,
                            region_shape: tuple = None,
                            scale: tuple = None,
                            units: tuple = None,
                            client: Client = None,
                            parallelize_over_regions = True,
                            func: Callable = otsu,
                            **func_params
                            ) -> zarr.Group:

    region_slices = get_regions(arr.shape, region_shape, as_slices = True)

    gr = zarr.open_group(output_path, mode='a')
    dataset = create_zarr_array(gr,
                                array_name = '0',
                                shape = arr.shape,
                                chunks = chunks,
                                dtype = arr.dtype,
                                overwrite = True
                                )

    meta = Multimeta()
    meta.parse_axes(axis_order='tczyx',
                    unit_list = units
                    )
    meta.add_dataset(path = '0',
                     scale = scale
                     )
    meta.to_ngff(gr)

    image_regions = [arr[region_slice] for region_slice in region_slices]
    # processed_regions = image_regions
    processed_regions = [func(reg, **func_params) for reg in image_regions]

    if client is not None:
        client.scatter(region_slices)
        client.scatter(image_regions)

    if not parallelize_over_regions:
        write_regions(image_regions = processed_regions,
                              region_slices = region_slices,
                              dataset = dataset,
                              client = client)
    else:
        write_regions_sequential(image_regions = processed_regions,
                                  region_slices = region_slices,
                                  dataset = dataset,
                                  client = client)
    return gr

def cellpose_seg(img,
                 model,
                 lower = 0.0,
                 upper = 800.0,
                 cell_diameter = 30.0):

    img = normalize_custom(img, lower = lower, upper = upper)
    masks, _, _ = model.eval(img,
                                channels = [0,0],
                                diameter = cell_diameter,
                                flow_threshold = 0.4,
                                cellprob_threshold = 0.0,
                                stitch_threshold = 0.5,
                                do_3D = False,
                                normalize = False,
                                )
    return masks

# if __name__ == '__main__':

chunks = (1, 1, 48, 128, 128)
region_shape = (1, 1, 91, 317, 325)
scale = (600, 1, 2, 0.406, 0.406)
units = ('s', 'Channel', 'µm', 'µm', 'µm')
# psf = gaussian_psf((1, 1, 12, 16, 16), (1, 1, 6, 8, 8), (1, 1, 12, 16, 16))
# psf = da.from_array(psf, chunks = chunks)

block_size = (1, 1, 5, 9, 9)

n_jobs = 4
threads_per_worker = 2
memory_limit = '8GB'

input_tiff_path_mg = f"/path/mg/*"
input_tiff_path_h2b = f"/path/H2B/*"

output_zarr_path = f"/path/concat_seg.zarr"

model_nuc = models.CellposeModel(gpu=True, model_type="CP")
model_cell = models.CellposeModel(gpu=True, model_type="CPx")

t0 = time.time()

paths_mg = sorted(glob.glob(input_tiff_path_mg))
paths_h2b = sorted(glob.glob(input_tiff_path_h2b))



# imgs_mg = [read_image(path) for path in paths_mg]
# imgs_h2b = [read_image(path) for path in paths_h2b]
#
# ### Concatenate collections into a single dask array
# mg_merged = da.concatenate(imgs_mg, axis=0)  # concatenate along the time dimension
# h2b_merged = da.concatenate(imgs_h2b, axis=0)  # concatenate along the time dimension
# imgs_merged = da.concatenate((mg_merged, h2b_merged), axis=1)  # concatenate along the channel dimension

# processed_img = da.concatenate([otsu(img, return_thresholded=True) for img in imgs_merged], axis=0)


with LocalCluster(processes=True,
                  nanny=True,
                  n_workers=n_jobs,
                  threads_per_worker=threads_per_worker,
                  memory_limit=memory_limit) as cluster:
    cluster.scale(n_jobs)
    with Client(cluster,
                heartbeat_interval="120s",
                timeout="600s",
                ) as client:

        ### Read image collections
        imgs_mg = [read_image(path) for path in paths_mg]
        imgs_h2b = [read_image(path) for path in paths_h2b]

        ### Concatenate collections into a single dask array
        mg_merged = da.concatenate(imgs_mg, axis = 0) # concatenate along the time dimension
        h2b_merged = da.concatenate(imgs_h2b, axis = 0) # concatenate along the time dimension
        imgs_merged = da.concatenate((mg_merged, h2b_merged), axis = 1) # concatenate along the channel dimension

        ### Process merged images
        processed_img = imgs_merged
        # processed_img = ndfilters.threshold_local(imgs_merged, block_size=block_size, method='mean')
        # processed_img = ndfilters.gaussian_filter(imgs_merged, sigma = (0.4, 0.4, 1, 1, 1))
        # filtered = ndfilters.uniform_filter(imgs_merged, size = block_size)
        # processed_img = imgs_merged > filtered
        # processed_mg = da.concatenate([utils.mean_threshold(img, return_thresholded=True) for img in imgs_mg], axis = 0)
        # processed_h2b = da.concatenate([utils.mean_threshold(img, return_thresholded=True) for img in imgs_h2b], axis = 0)
        # processed_mg = da.concatenate([utils.otsu(img, bincount = 9, return_thresholded=True) for img in imgs_mg], axis = 0)
        # processed_h2b = da.concatenate([utils.otsu(img, bincount = 9, return_thresholded=True) for img in imgs_h2b], axis = 0)
        # processed_img = da.concatenate((processed_mg, processed_h2b), axis = 1) # concatenate along the channel dimension

        process_and_save_to_ngff(processed_img,
                output_path = output_zarr_path,
                region_shape = region_shape,
                scale = scale,
                units = units,
                client = client,
                parallelize_over_regions=False,
                func = cellpose_seg,
                model = model_nuc
                )