Add option to cache background mask in object measures and add script for processing all ChReef cochleae

Author: constantinpape

flamingo_tools/measurements.py

@@ -7,6 +7,7 @@
 import numpy as np
 import pandas as pd
 import trimesh
+from elf.io import open_file
 from elf.wrapper.resized_volume import ResizedVolume
 from nifty.tools import blocking
 from skimage.measure import marching_cubes, regionprops_table
@@ -261,6 +262,7 @@ def compute_object_measures_impl(
 
     # For debugging.
     # measure_function(seg_ids[0])
+    # breakpoint()
 
     with futures.ThreadPoolExecutor(n_threads) as pool:
         measures = list(tqdm(
@@ -352,6 +354,7 @@ def compute_sgn_background_mask(
     threshold_percentile: float = 35.0,
     scale_factor: Tuple[int, int, int] = (16, 16, 16),
     n_threads: Optional[int] = None,
+    cache_path: Optional[str] = None,
 ) -> np.typing.ArrayLike:
     """Compute the background mask for intensity measurements in the SGN segmentation.
 
@@ -368,6 +371,7 @@ def compute_sgn_background_mask(
         threshold_percentile: The percentile threshold for separating foreground and background in the PV signal.
         scale_factor: The scale factor for internally downsampling the mask.
         n_threads: The number of threads for parallelizing the computation.
+        cache_path: Optional path to save the downscaled background mask to zarr.
 
     Returns:
         The mask for determining the background values.
@@ -376,6 +380,13 @@ def compute_sgn_background_mask(
     segmentation = read_image_data(segmentation_path, segmentation_key)
     assert image.shape == segmentation.shape
 
+    if cache_path is not None and os.path.exists(cache_path):
+        with open_file(cache_path, "r") as f:
+            if "mask" in f:
+                low_res_mask = f["mask"][:]
+                mask = ResizedVolume(low_res_mask, shape=image.shape, order=0)
+                return mask
+
     original_shape = image.shape
     downsampled_shape = tuple(int(np.round(sh / sf)) for sh, sf in zip(original_shape, scale_factor))
 
@@ -406,10 +417,16 @@ def _compute_block(block_id):
         low_res_mask[bb] = this_mask
 
     n_threads = mp.cpu_count() if n_threads is None else n_threads
+    randomized_blocks = np.arange(0, n_blocks)
+    np.random.shuffle(randomized_blocks)
     with futures.ThreadPoolExecutor(n_threads) as tp:
         list(tqdm(
-            tp.map(_compute_block, range(n_blocks)), total=n_blocks, desc="Compute background mask"
+            tp.map(_compute_block, randomized_blocks), total=n_blocks, desc="Compute background mask"
         ))
 
+    if cache_path is not None:
+        with open_file(cache_path, "a") as f:
+            f.create_dataset("mask", data=low_res_mask, chunks=(64, 64, 64))
+
     mask = ResizedVolume(low_res_mask, shape=original_shape, order=0)
     return mask

reproducibility/object_measures/ChReef_MLR143R.json

@@ -0,0 +1,13 @@
+[
+    {
+        "cochlea": "M_LR_000143_R",
+        "image_channel": [
+            "GFP"
+        ],
+        "segmentation_channel": "SGN_v2",
+        "background_mask": "yes",
+        "component_list": [
+            1
+        ]
+    }
+]

reproducibility/object_measures/process_all_object_measures.py

@@ -0,0 +1,74 @@
+import json
+import os
+import subprocess
+import zarr
+
+import flamingo_tools.s3_utils as s3_utils
+
+OUTPUT_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/mobie_project/cochlea-lightsheet/tables/measurements/"  # noqa
+JSON_ROOT = "/user/pape41/u12086/Work/my_projects/flamingo-tools/reproducibility/object_measures"
+COCHLEAE = [
+    "M_LR_000143_L",
+    "M_LR_000144_L",
+    "M_LR_000145_L",
+    "M_LR_000153_L",
+    "M_LR_000155_L",
+    "M_LR_000189_L",
+    "M_LR_000143_R",
+    "M_LR_000144_R",
+    "M_LR_000145_R",
+    "M_LR_000153_R",
+    "M_LR_000155_R",
+    "M_LR_000189_R",
+]
+
+
+def process_cochlea(cochlea, start_slurm):
+    short_name = cochlea.replace("_", "").replace("0", "")
+
+    # Check if this cochlea has been processed already.
+    output_name = cochlea.replace("_", "-")
+    output_path = os.path.join(OUTPUT_ROOT, f"{output_name}_GFP_SGN-v2_object-measures.tsv")
+    if os.path.exists(output_path):
+        print(cochlea, "has been processed already.")
+        return
+
+    # Check if the raw data for this cochlea is accessible.
+    img_name = f"{cochlea}/images/ome-zarr/GFP.ome.zarr"
+    img_path, _ = s3_utils.get_s3_path(img_name)
+    try:
+        zarr.open(img_path, mode="r")
+    except Exception:
+        print("The data for", cochlea, "at", img_name, "does not exist.")
+        return
+
+    # Then generate the json file if it does not yet exist.
+    template_path = os.path.join(JSON_ROOT, "ChReef_MLR143L.json")
+    with open(template_path, "r") as f:
+        json_template = json.load(f)
+
+    json_path = os.path.join(JSON_ROOT, f"ChReef_{short_name}.json")
+    if not os.path.exists(json_path):
+        print("Write json to", json_path)
+        # TODO: We may need to replace the component list for some.
+        json_template[0]["cochlea"] = cochlea
+        with open(json_path, "w") as f:
+            json.dump(json_template, f, indent=4)
+
+    print(cochlea, "is not yet processed")
+    # Then start the slurm job.
+    if not start_slurm:
+        return
+
+    print("Submit slurm job for", cochlea)
+    subprocess.run(["sbatch", "slurm_template.sbatch", json_path, OUTPUT_ROOT])
+
+
+def main():
+    start_slurm = False
+    for cochlea in COCHLEAE:
+        process_cochlea(cochlea, start_slurm)
+
+
+if __name__ == "__main__":
+    main()

reproducibility/object_measures/repro_object_measures.py

@@ -1,6 +1,7 @@
 import argparse
 import json
 import os
+from multiprocessing import cpu_count
 from typing import Optional
 
 import flamingo_tools.s3_utils as s3_utils
@@ -49,6 +50,7 @@ def repro_object_measures(
             seg_path, fs = s3_utils.get_s3_path(seg_s3, bucket_name=s3_bucket_name,
                                                 service_endpoint=s3_service_endpoint, credential_file=s3_credentials)
 
+            n_threads = int(os.environ.get("SLURM_CPUS_ON_NODE", cpu_count()))
             if os.path.isfile(output_table_path) and not force_overwrite:
                 print(f"Skipping creation of {output_table_path}. File already exists.")
 
@@ -62,11 +64,14 @@ def repro_object_measures(
                     feature_set = "default_background_subtract"
                     dilation = 4
                     median_only = True
+                    mask_cache_path = os.path.join(output_dir, f"{cochlea_str}_{img_str}_{seg_str}_bg-mask.zarr")
                     bg_mask = compute_sgn_background_mask(
                         image_path=img_path,
                         segmentation_path=seg_path,
                         image_key=input_key,
                         segmentation_key=input_key,
+                        n_threads=n_threads,
+                        cache_path=mask_cache_path,
                     )
 
                 compute_object_measures(
@@ -82,6 +87,7 @@ def repro_object_measures(
                     dilation=dilation,
                     median_only=median_only,
                     background_mask=bg_mask,
+                    n_threads=n_threads,
                 )
 
 

reproducibility/object_measures/slurm_template.sbatch

@@ -0,0 +1,15 @@
+#!/bin/bash
+#SBATCH -t 24:00:00             # estimated time, adapt to your needs
+#SBATCH [email protected] # change this to your mailaddress
+#SBATCH --mail-type=FAIL        # send mail when job begins and ends
+
+#SBATCH -p standard96s:shared         # the partition
+#SBATCH -A nim00007
+
+#SBATCH -c 32
+#SBATCH --mem 256G
+
+PYTHON=/scratch-grete/usr/nimcpape/software/mamba/envs/sam/bin/python
+SCRIPT=/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/reproducibility/object_measures/repro_object_measures.py
+
+$PYTHON $SCRIPT --input $1 --output $2