image_labeling prototype (ref #64)

fractal/tasks/image_labeling.py

@@ -0,0 +1,75 @@
+"""
+Copyright 2022 (C) Friedrich Miescher Institute for Biomedical Research and
+University of Zurich
+
+Original authors:
+Tommaso Comparin <[email protected]>
+Marco Franzon <[email protected]>
+
+This file is part of Fractal and was originally developed by eXact lab S.r.l.
+<exact-lab.it> under contract with Liberali Lab from the Friedrich Miescher
+Institute for Biomedical Research and Pelkmans Lab from the University of
+Zurich.
+"""
+# import json
+import dask.array as da
+import zarr
+from cellpose import core
+from cellpose import models
+from ome_zarr.io import parse_url
+from ome_zarr.writer import write_labels
+
+
+def image_labeling(
+    zarrurl,
+):
+
+    """
+    FIXME
+    """
+
+    # Read number of levels from .zattrs of original zarr file
+    # with open(zarrurl + ".zattrs", "r") as inputjson:
+    #     zattrs = json.load(inputjson)
+    # num_levels = len(zattrs["multiscales"][0]["datasets"])
+
+    # Load some level and some channel #FIXME
+    dapi_dset = da.from_zarr(zarrurl + "/0")[0]
+    print("LOADED ZARR")
+    print(dapi_dset.shape)
+    print()
+
+    use_gpu = core.use_gpu()
+    model = models.Cellpose(gpu=use_gpu, model_type="nuclei")
+    mask, flows, styles, diams = model.eval(
+        dapi_dset, channels=[0, 0], do_3D=True, net_avg=False, augment=False
+    )
+
+    print(type(mask), mask.shape)
+    print(type(flows), len(flows), len(flows[0]))
+    print(type(styles), styles.shape)
+    print(type(diams), diams)
+
+    store = parse_url(zarrurl, mode="w").store
+    root = zarr.group(store=store)
+    label_name = "label_image"
+    label_axes = [
+        "z",
+        "y",
+        "x",
+    ]  # could change if e.g. a 2D image was processed => can we infer this?
+    write_labels(mask, group=root, name=label_name, axes=label_axes)
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser(prog="image_labeling.py")
+    parser.add_argument(
+        "-z", "--zarrurl", help="zarr url, at the FOV level", required=True
+    )
+
+    args = parser.parse_args()
+    image_labeling(
+        args.zarrurl,
+    )