Implement mobie export and combined segmentation workflow

constantinpape · constantinpape · commit 0c628da51cc0 · 2024-09-08T12:28:48.000+02:00
diff --git a/flamingo_tools/data_conversion.py b/flamingo_tools/data_conversion.py
@@ -134,10 +134,16 @@ def flamingo_filename_parser(file_path, name_mapping):
     illumination_mapping = name_mapping.get("illumination", {})
     attributes["illumination"] = {"id": illumination, "name": illumination_mapping.get(illumination, str(illumination))}
 
+    # Extract D. TODO what is this?
+    match = re.search(r'_D(\d+)_', filename)
+    D = int(match.group(1)) if match else 0
+    D_mapping = name_mapping.get("D", {})
+    attributes["D"] = {"id": D, "name": D_mapping.get(D, str(D))}
+
     # BDV also supports an angle attribute, but it does not seem to be stored in the filename
     # "angle": {"id": 0, "name": "0"}
 
-    attribute_id = f"c{channel}-t{tile}-i{illumination}"
+    attribute_id = f"c{channel}-t{tile}-i{illumination}-d{D}"
     return timepoint, attributes, attribute_id
 
 
@@ -242,13 +248,13 @@ def convert_lightsheet_to_bdv(
         else:  # We have metadata and read it.
             resolution, unit, tile_transformation = read_metadata_flamingo(metadata_file, offset)
 
+        print(f"Converting tp={timepoint}, channel={attributes['channel']}, tile={attributes['tile']}")
         try:
             data = tifffile.memmap(file_path, mode="r")
         except ValueError:
             print(f"Could not memmap the data from {file_path}. Fall back to load it into memory.")
             data = tifffile.imread(file_path)
 
-        print(f"Converting tp={timepoint}, channel={attributes['channel']}, tile={attributes['tile']}")
         if scale_factors is None:
             scale_factors = derive_scale_factors(data.shape)
 
diff --git a/flamingo_tools/mobie.py b/flamingo_tools/mobie.py
@@ -0,0 +1,77 @@
+import os
+import tempfile
+from typing import Tuple
+
+from mobie import add_bdv_image, add_segmentation
+from mobie.metadata.dataset_metadata import read_dataset_metadata
+
+
+# TODO refactor to mobie utils
+def _source_exists(mobie_project, mobie_dataset, source_name):
+    dataset_folder = os.path.join(mobie_project, mobie_dataset)
+    metadata = read_dataset_metadata(dataset_folder)
+    sources = metadata.get("sources", {})
+    return source_name in sources
+
+
+def add_raw_to_mobie(
+    mobie_project: str,
+    mobie_dataset: str,
+    source_name: str,
+    xml_path: str,
+    skip_existing: bool = True,
+    setup_id: int = 0,
+):
+    """
+    """
+    # Check if we have converted this data already.
+    have_source = _source_exists(mobie_project, mobie_dataset, source_name)
+    if have_source and skip_existing:
+        print(f"Source {source_name} already exists in {mobie_project}:{mobie_dataset}.")
+        print("Conversion to mobie will be skipped.")
+        return
+    elif have_source:
+        raise NotImplementedError
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        add_bdv_image(
+            xml_path=xml_path,
+            root=mobie_project,
+            dataset_name=mobie_dataset,
+            image_name=source_name,
+            tmp_folder=tmpdir,
+            file_format="bdv.n5",
+            setup_ids=[setup_id],
+        )
+
+
+def add_segmentation_to_mobie(
+    mobie_project: str,
+    mobie_dataset: str,
+    source_name: str,
+    segmentation_path: str,
+    segmentation_key: str,
+    resolution: Tuple[int, int, int],
+    unit: str,
+    scale_factors: Tuple[Tuple[int, int, int]],
+    chunks: Tuple[int, int, int],
+    skip_existing: bool = True,
+):
+    # Check if we have converted this data already.
+    have_source = _source_exists(mobie_project, mobie_dataset, source_name)
+    if have_source and skip_existing:
+        print(f"Source {source_name} already exists in {mobie_project}:{mobie_dataset}.")
+        print("Conversion to mobie will be skipped.")
+        return
+    elif have_source:
+        raise NotImplementedError
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        add_segmentation(
+            input_path=segmentation_path, input_key=segmentation_key,
+            root=mobie_project, dataset_name=mobie_dataset,
+            segmentation_name=source_name,
+            resolution=resolution, scale_factors=scale_factors,
+            chunks=chunks, file_format="bdv.n5",
+            tmp_folder=tmpdir
+        )
diff --git a/scripts/data_transfer/README.md b/scripts/data_transfer/README.md
@@ -18,8 +18,7 @@ $ smbclient \\\\wfs-medizin.top.gwdg.de\\ukon-all\$\\ukon100 -U GWDG\\pape41"
     - mget *
 - Copy this to HLRN by logging into it and running
 ```
-  $ rsync -e "ssh -i ~/.ssh/id_rsa_hlrn" -avz pape41@login-mdc.hpc.gwdg.de:/scratch1/projects/cca/data/mose
-r/<NAME> /mnt/lustre-emmy-hdd/projects/nim00007/data/moser/lightsheet/volumes/<NAME>
+  $ rsync -e "ssh -i ~/.ssh/id_rsa_hlrn" -avz pape41@login-mdc.hpc.gwdg.de:/scratch1/projects/cca/data/moser/<NAME> /mnt/lustre-emmy-hdd/projects/nim00007/data/moser/lightsheet/volumes/<NAME>
 ```
 - Remove on SCC
 
@@ -28,9 +27,11 @@ r/<NAME> /mnt/lustre-emmy-hdd/projects/nim00007/data/moser/lightsheet/volumes/<N
 - UKON100\archiv\imaging\Lightsheet\Huiskengroup_CTLSM\2024\M171_2R_converted_n5
     - unclear what the converted data is
 - UKON100\archiv\imaging\Lightsheet\Huiskengroup_CTLSM\2024\155_1L_converted_n5\BDVexport.n5
-    - Copied to SCC, need to rsync.
+    - Copied to SCC, rsync in progress.
 - UKON100\archiv\imaging\Lightsheet\Huiskengroup_CTLSM\2024\MLR151_2R_converted_n5
+    - unclear what the converted data is
 - UKON100\archiv\imaging\Lightsheet\Huiskengroup_CTLSM\2024\G11_1L_converted_n5
+    - unclear what the converted data is
 
 ## Improvements
 
diff --git a/scripts/prediction/segmentation_workflow.py b/scripts/prediction/segmentation_workflow.py
@@ -0,0 +1,152 @@
+import argparse
+import os
+from shutil import rmtree
+
+import pybdv.metadata as bdv_metadata
+import torch
+import z5py
+
+from flamingo_tools.segmentation import run_unet_prediction, filter_isolated_objects
+from flamingo_tools.mobie import add_raw_to_mobie, add_segmentation_to_mobie
+
+MOBIE_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/moser/lightsheet/mobie"
+
+
+def postprocess_seg(output_folder):
+    print("Run segmentation postprocessing ...")
+    seg_path = os.path.join(output_folder, "segmentation.zarr")
+    seg_key = "segmentation"
+
+    with z5py.File(seg_path, "r") as f:
+        segmentation = f[seg_key][:]
+
+    seg_filtered, n_pre, n_post = filter_isolated_objects(segmentation)
+
+    with z5py.File(seg_path, "a") as f:
+        chunks = f[seg_key].chunks
+        f.create_dataset(
+            "segmentation_postprocessed", data=seg_filtered, compression="gzip",
+            chunks=chunks, dtype=seg_filtered.dtype
+        )
+
+
+def export_to_mobie(xml_path, segmentation_folder, scale, mobie_dataset, chunks):
+    # Add to mobie:
+
+    # - raw data (if not yet present)
+    add_raw_to_mobie(
+        mobie_project=MOBIE_ROOT,
+        mobie_dataset=mobie_dataset,
+        source_name="pv-channel",
+        xml_path=xml_path,
+        setup_id=0,
+    )
+
+    # TODO enable passing extra channel names
+    # - additional channels
+    setup_ids = bdv_metadata.get_setup_ids(xml_path)
+    if len(setup_ids) > 1:
+        extra_channel_names = ["gfp_channel", "myo_channel"]
+        for i, setup_id in enumerate(setup_ids[1:]):
+            add_raw_to_mobie(
+                mobie_project=MOBIE_ROOT,
+                mobie_dataset=mobie_dataset,
+                source_name=extra_channel_names[i],
+                xml_path=xml_path,
+                setup_id=setup_id
+            )
+
+    # - segmentation and post-processed segmentation
+    seg_path = os.path.join(segmentation_folder, "segmentation.zarr")
+    seg_resolution = bdv_metadata.get_resolution(xml_path, setup_id=0)
+    if scale == 1:
+        seg_resolution = [2 * res for res in seg_resolution]
+    unit = bdv_metadata.get_unit(xml_path, setup_id=0)
+
+    seg_key = "segmentation"
+    seg_name = "nuclei_fullscale" if scale == 0 else "nuclei_downscaled"
+    add_segmentation_to_mobie(
+        mobie_project=MOBIE_ROOT,
+        mobie_dataset=mobie_dataset,
+        source_name=seg_name,
+        segmentation_path=seg_path,
+        segmentation_key=seg_key,
+        resolution=seg_resolution,
+        unit=unit,
+        scale_factors=4*[[2, 2, 2]],
+        chunks=chunks,
+    )
+
+    seg_key = "segmentation_postprocessed"
+    seg_name += "_postprocessed"
+    add_segmentation_to_mobie(
+        mobie_project=MOBIE_ROOT,
+        mobie_dataset=mobie_dataset,
+        source_name=seg_name,
+        segmentation_path=seg_path,
+        segmentation_key=seg_key,
+        resolution=seg_resolution,
+        unit=unit,
+        scale_factors=4*[[2, 2, 2]],
+        chunks=chunks,
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input", required=True)
+    parser.add_argument("-o", "--output_folder", required=True)
+    parser.add_argument("-s", "--scale", required=True, type=int)
+    parser.add_argument("-m", "--mobie_dataset", required=True)
+    parser.add_argument("--model")
+
+    args = parser.parse_args()
+
+    scale = args.scale
+    if scale == 0:
+        min_size = 1000
+    elif scale == 1:
+        min_size = 250
+    else:
+        raise ValueError
+
+    xml_path = args.input
+    assert os.path.splitext(xml_path)[1] == ".xml"
+    input_path = bdv_metadata.get_data_path(xml_path, return_absolute_path=True)
+
+    # TODO need to make sure that PV is always setup 0
+    input_key = f"setup0/timepoint0/s{scale}"
+
+    have_cuda = torch.cuda.is_available()
+    chunks = z5py.File(input_path, "r")[input_key].chunks
+    block_shape = tuple([2 * ch for ch in chunks]) if have_cuda else tuple(chunks)
+    halo = (16, 64, 64) if have_cuda else (8, 32, 32)
+
+    if args.model is not None:
+        model = args.model
+    else:
+        if scale == 0:
+            model = "../training/checkpoints/cochlea_distance_unet"
+        else:
+            model = "../training/checkpoints/cochlea_distance_unet-train-downsampled"
+
+    run_unet_prediction(
+        input_path, input_key, args.output_folder, model,
+        scale=None, min_size=min_size,
+        block_shape=block_shape, halo=halo,
+    )
+
+    postprocess_seg(args.output_folder)
+
+    export_to_mobie(xml_path, args.output_folder, scale, args.mobie_dataset, chunks)
+
+    # clean up: remove segmentation folders
+    print("Cleaning up intermediate segmentation results")
+    print("This may take a while, but everything else is done.")
+    print("You can check the results in the MoBIE project already at:")
+    print(f"{MOBIE_ROOT}:{args.mobie_dataset}")
+    rmtree(args.output_folder)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/prediction/upload_to_s3.py b/scripts/prediction/upload_to_s3.py
@@ -0,0 +1 @@
+# TODO find suitable bucket
