Update La-Vision processing

Author: constantinpape

scripts/la-vision/import_mouse_lsm.py

@@ -2,6 +2,7 @@
 
 import imageio.v3 as imageio
 from mobie import add_image
+from mobie.metadata import read_dataset_metadata
 
 INPUT_ROOT = "/mnt/ceph-hdd/cold/nim00007/cochlea-lightsheet/keppeler-et-al/mouse"
 MOBIE_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/mobie_project/cochlea-lightsheet"
@@ -17,6 +18,12 @@ def add_mouse_lsm():
     chunks = (96, 96, 96)
 
     for channel_id, channel_name in enumerate(channel_names, 1):
+        mobie_ds_folder = os.path.join(MOBIE_ROOT, DS_NAME)
+        ds_metadata = read_dataset_metadata(mobie_ds_folder)
+        if channel_name in ds_metadata["sources"]:
+            print(channel_name, "is already in MoBIE")
+            continue
+
         input_path = os.path.join(INPUT_ROOT, f"Mouse_cochlea_04_LSFM_ch{channel_id}_raw.tif")
         print("Load image data ...")
         input_data = imageio.imread(input_path)

scripts/la-vision/predict_blocks.py

@@ -3,31 +3,50 @@
 from glob import glob
 
 import imageio.v3 as imageio
+import numpy as np
+
+from skimage.segmentation import watershed
+from skimage.measure import label
 from torch_em.util import load_model
 from torch_em.util.prediction import predict_with_halo
 
-INPUT_ROOT = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/more-annotations/LA_VISION_M04"  # noqa
 
+def _get_files(sgn=True):
+    input_root = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/more-annotations/LA_VISION_M04"  # noqa
+    input_root2 = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/more-annotations/LA_VISION_M04_2"  # noqa
+    input_root3 = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/more-annotations/LA_VISION_Mar05"  # noqa
+
+    input_root_ihc = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/more-annotations/LA_VISION_Mar05-ihc"  # noqa
+
+    if sgn:
+        input_files = glob(os.path.join(input_root, "*.tif")) +\
+            glob(os.path.join(input_root2, "*.tif")) +\
+            glob(os.path.join(input_root3, "*.tif"))
+    else:
+        input_files = glob(os.path.join(input_root_ihc, "*.tif"))
+
+    return input_files
 
-def predict_blocks(model_path, name):
+
+def predict_blocks(model_path, name, sgn=True):
     output_folder = os.path.join("./predictions", name)
     os.makedirs(output_folder, exist_ok=True)
 
-    input_blocks = glob(os.path.join(INPUT_ROOT, "*.tif"))
+    input_blocks = _get_files(sgn)
 
-    model = load_model(model_path)
+    model = None
     for path in input_blocks:
+        out_path = os.path.join(output_folder, os.path.basename(path))
+        if os.path.exists(out_path):
+            continue
+        if model is None:
+            model = load_model(model_path)
         data = imageio.imread(path)
         pred = predict_with_halo(data, model, gpu_ids=[0], block_shape=[64, 128, 128], halo=[8, 32, 32])
-        out_path = os.path.join(output_folder, os.path.basename(path))
         imageio.imwrite(out_path, pred, compression="zlib")
 
 
 def _segment_impl(pred, dist_threshold=0.5):
-    import numpy as np
-    from skimage.segmentation import watershed
-    from skimage.measure import label
-
     fg, center_dist, boundary_dist = pred
     mask = fg > 0.5
 
@@ -37,35 +56,42 @@ def _segment_impl(pred, dist_threshold=0.5):
     return seg
 
 
-def check_segmentation(name):
+def check_segmentation(name, sgn):
     import napari
 
-    input_blocks = sorted(glob(os.path.join(INPUT_ROOT, "*.tif")))
+    input_files = _get_files(sgn)
 
     output_folder = os.path.join("./predictions", name)
-    pred = sorted(glob(os.path.join(output_folder, "*.tif")))
 
-    for path, pred_path in zip(input_blocks, pred):
+    for path in input_files:
         image = imageio.imread(path)
+        pred_path = os.path.join(output_folder, os.path.basename(path))
         pred = imageio.imread(pred_path)
-        seg = _segment_impl(pred)
+        if sgn:
+            seg = _segment_impl(pred)
+        else:
+            seg = label(pred[0] > 0.5)
         v = napari.Viewer()
         v.add_image(image)
         v.add_image(pred)
         v.add_labels(seg)
+        v.title = os.path.basename(path)
         napari.run()
 
 
 # Model path for original training on low-res SGNs:
 # /mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/training/checkpoints/cochlea_distance_unet_low-res-sgn  # noqa
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", "-m", required=True)
     parser.add_argument("--name", "-n", required=True)
+    parser.add_argument("--model_path", "-m")
+    parser.add_argument("--check", action="store_true")
+    parser.add_argument("--ihc", action="store_true")
     args = parser.parse_args()
 
-    # predict_blocks(args.model_path, args.name)
-    check_segmentation(args.name)
+    predict_blocks(args.model_path, args.name, sgn=not args.ihc)
+    if args.check:
+        check_segmentation(args.name, sgn=not args.ihc)
 
 
 if __name__ == "__main__":

scripts/la-vision/segment_sgns.py

@@ -7,9 +7,9 @@
 from mobie import add_segmentation
 from mobie.metadata import add_remote_project_metadata
 
-MODEL_PATH = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/training/checkpoints/cochlea_distance_unet_low-res-sgn"  # noqa
+MODEL_PATH = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/training/checkpoints/cochlea_distance_unet_low-res-sgn-v2"  # noqa
 RESOLUTION = (3.0, 1.887779, 1.887779)
-SEG_NAME = "SGN_LOWRES"
+SEG_NAME = "SGN_LOWRES-v2"
 
 
 def segment_sgns(input_path, input_key, output_folder):
@@ -56,6 +56,7 @@ def upload_to_s3(mobie_dir, dataset_name):
 
     add_remote_project_metadata(mobie_dir, bucket_name, service_endpoint)
 
+    # run(["module", "load", "rclone"])
     run(["rclone", "--progress", "copyto",
          f"{mobie_dir}/{dataset_name}/dataset.json",
          f"cochlea-lightsheet:cochlea-lightsheet/{dataset_name}/dataset.json"])
@@ -92,7 +93,7 @@ def segment_Mar05():
 
 
 def main():
-    # segment_M04()
+    segment_M04()
     segment_Mar05()
 
 

scripts/more-annotations/extract_sgn_annotations.py

@@ -96,21 +96,48 @@ def downscale_segmentation():
 
 # Note: consider different normalization strategy for these cochleae and normalize by local intensity
 # rather than by global values.
-
 # Also double check empty positions again and make sure they don't contain SGNs
-# Additional positions for LaVision annotations:
-# {"position":[2031.0655170248258,1925.206039671767,249.14546086048554],"timepoint":0}
-# {"position":[2378.3720460599393,2105.471228531872,303.9285928812524],"timepoint":0}
-# {"position":[1619.3251178227529,3444.7351705689553,271.2360278843609],"timepoint":0}
-# {"position":[2358.2784398426843,1503.2211953830192,762.7325586759833],"timepoint":0}
+def download_lavision_crops2():
+    input_key = "s0"
+    output_key = None
 
-# Position in Marmoset:
-# {"position":[2462.7875134103206,2818.067344942212,1177.1380214828991],"timepoint":0}
+    # Additional positions for LaVision annotations:
+    input_path = "LaVision-M04/images/ome-zarr/PV.ome.zarr"
+    new_positions_m04 = [
+        [2031.0655170248258, 1925.206039671767, 249.14546086048554],
+        [2378.3720460599393, 2105.471228531872, 303.9285928812524],
+        [1619.3251178227529, 3444.7351705689553, 271.2360278843609],
+        [2358.2784398426843, 1503.2211953830192, 762.7325586759833],
+    ]
+
+    output_folder = "./LA_VISION_M04_2"
+    os.makedirs(output_folder, exist_ok=True)
+    for pos in new_positions_m04:
+        halo = [128, 128, 32]
+        extract_block(
+            input_path, pos, output_folder, input_key, output_key, RESOLUTION_LA_VISION, halo,
+            tif=True, s3=True,
+        )
+
+    # Position in Marmoset:
+    new_positions_mar05 = [
+        [2462.7875134103206, 2818.067344942212, 1177.1380214828991]
+    ]
+    input_path = "LaVision-Mar05/images/ome-zarr/PV.ome.zarr"
+    output_folder = "./LA_VISION_Mar05"
+    os.makedirs(output_folder, exist_ok=True)
+    for pos in new_positions_mar05:
+        halo = [128, 128, 32]
+        extract_block(
+            input_path, pos, output_folder, input_key, output_key, RESOLUTION_LA_VISION, halo,
+            tif=True, s3=True,
+        )
 
 
 def main():
     # download_lavision_crops()
-    downscale_segmentation()
+    # downscale_segmentation()
+    download_lavision_crops2()
 
 
 if __name__ == "__main__":