Merge pull request #41 from computational-cell-analytics/add_baselines

The baseline methods were evaluated for the initial submission of the paper.
The branch has been brought to a satisfactory state and can now be merged into the master branch.

Author: schilling40

flamingo_tools/validation.py

@@ -217,7 +217,7 @@ def compute_matches_for_annotated_slice(
         A dictionary with keys 'tp_objects', 'tp_annotations' 'fp' and 'fn', mapping to the respective ids.
     """
     assert segmentation.ndim in (2, 3)
-    coordinates = ["axis-0", "axis-1"] if segmentation.ndim == 2 else ["axis-0", "axis-1", "axis-2"]
+    coordinates = ["axis-1", "axis-2"] if segmentation.ndim == 2 else ["axis-0", "axis-1", "axis-2"]
     segmentation_ids = np.unique(segmentation)[1:]
 
     # Crop to the minimal enclosing bounding box of points and segmented objects.

scripts/baselines/NIS3D_apply.py

@@ -0,0 +1,46 @@
+import os
+import sys
+
+script_dir = "/user/schilling40/u15000/flamingo-tools/scripts/prediction"
+sys.path.append(script_dir)
+
+import run_prediction_distance_unet
+
+checkpoint_dir = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/trained_models/nucleus"
+model_name = "NIS3D_supervised_2025-07-17"
+model_dir = os.path.join(checkpoint_dir, model_name)
+checkpoint = os.path.join(checkpoint_dir, model_name, "best.pt")
+
+cochlea_dir = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet"
+
+image_dir = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/nucleus/2025-07_NIS3D/test"
+
+out_dir = os.path.join(cochlea_dir, "predictions", "val_nucleus", "distance_unet_NIS3D")  # /distance_unet
+
+boundary_distance_threshold = 0.5
+seg_class = "ihc"
+
+block_shape = (128, 128, 128)
+halo = (16, 32, 32)
+
+block_shape_str = ",".join([str(b) for b in block_shape])
+halo_str = ",".join([str(h) for h in halo])
+
+images = [entry.path for entry in os.scandir(image_dir) if entry.is_file() and "iitest.tif" in entry.path]
+
+for image in images:
+    sys.argv = [
+        os.path.join(script_dir, "run_prediction_distance_unet.py"),
+        f"--input={image}",
+        f"--output_folder={out_dir}",
+        f"--model={model_dir}",
+        f"--block_shape=[{block_shape_str}]",
+        f"--halo=[{halo_str}]",
+        "--memory",
+        "--time",
+        "--no_masking",
+        f"--seg_class={seg_class}",
+        f"--boundary_distance_threshold={boundary_distance_threshold}"
+    ]
+
+    run_prediction_distance_unet.main()

scripts/baselines/NIS3D_eval.py

@@ -0,0 +1,160 @@
+import argparse
+import json
+import multiprocessing as mp
+import os
+from concurrent import futures
+from typing import List
+
+import numpy as np
+import tifffile
+from tqdm import tqdm
+
+GT_DIR = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/nucleus/2025-07_NIS3D/test"
+PRED_DIR = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/predictions/val_nucleus/distance_unet_NIS3D"
+
+
+def find_overlapping_masks(
+    arr_base: np.ndarray,
+    arr_ref: np.ndarray,
+    label_id_base: int,
+    min_overlap: float = 0.5,
+) -> List[int]:
+    """Find masks of segmentation, which have an overlap with undefined mask greater than 0.5.
+    """
+    labels_undefined_mask = []
+    arr_base_undefined = arr_base == label_id_base
+
+    # iterate through segmentation ids in reference mask
+    ref_ids = list(np.unique(arr_ref)[1:])
+    for ref_id in ref_ids:
+        arr_ref_instance = arr_ref == ref_id
+
+        intersection = np.logical_and(arr_ref_instance, arr_base_undefined)
+        overlap_ratio = np.sum(intersection) / np.sum(arr_ref_instance)
+        if overlap_ratio >= min_overlap:
+            labels_undefined_mask.append(ref_id)
+
+    return labels_undefined_mask
+
+
+def find_matching_masks(arr_gt, arr_ref, out_path, labels_undefined_mask=[]):
+    """For each instance in the reference array, the corresponding mask of the ground truth array,
+    which has the biggest overlap, is identified.
+
+    Args:
+        arr_gt:
+        arr_ref:
+        out_path: Output path for saving dictionary.
+        labels_undefined_mask: Labels of the reference array to exclude.
+    """
+    seg_ids_ref = [int(i) for i in np.unique(arr_ref)[1:]]
+    print(f"total number of segmentation masks: {len(seg_ids_ref)}")
+    seg_ids_ref = [s for s in seg_ids_ref if s not in labels_undefined_mask]
+    print(f"number of segmentation masks after filtering undefined masks: {len(seg_ids_ref)}")
+
+    def compute_overlap(ref_id):
+        """Identify ID of segmentation mask with biggest overlap.
+        Return matched IDs and overlap.
+        """
+        arr_ref_instance = arr_ref == ref_id
+
+        seg_ids_gt = np.unique(arr_gt[arr_ref_instance])[1:]
+
+        max_overlap = 0
+        gt_id_match = None
+
+        for gt_id in seg_ids_gt:
+            arr_gt_instance = arr_gt == gt_id
+
+            intersection = np.logical_and(arr_ref_instance, arr_gt_instance)
+            overlap_ratio = np.sum(intersection) / np.sum(arr_ref_instance)
+            if overlap_ratio > max_overlap:
+                gt_id_match = int(gt_id.tolist())
+                max_overlap = np.max([max_overlap, overlap_ratio])
+
+        if gt_id_match is not None:
+            return {
+                "ref_id": ref_id,
+                "gt_id": gt_id_match,
+                "overlap": float(max_overlap.tolist())
+            }
+        else:
+            return None
+
+    n_threads = min(16, mp.cpu_count())
+    print(f"Parallelizing with {n_threads} Threads.")
+    with futures.ThreadPoolExecutor(n_threads) as pool:
+        results = list(tqdm(pool.map(compute_overlap, seg_ids_ref), total=len(seg_ids_ref)))
+
+    matching_masks = {r['ref_id']: r for r in results if r is not None}
+
+    with open(out_path, "w") as f:
+        json.dump(matching_masks, f, indent='\t', separators=(',', ': '))
+
+
+def filter_true_positives(output_folder, prefixes, force_overwrite):
+    """ Filter true positives from segmentation.
+    Segmentation instances and ground truth labels are filtered symmetrically.
+    The maximal overlap of each is computed and taken as a true positive if symmetric.
+    The instance ID, the reference ID, and the overlap are saved in dictionaries.
+
+    Args:
+        output_folder: Output folder for dictionaries.
+        prefixes: List of prefixes for evaluation. One or multiple of ["Drosophila", "MusMusculus", "Zebrafish"].
+        force_overwrite: Flag for forced overwrite of existing output files.
+    """
+    if "PRED_DIR" in globals():
+        pred_dir = PRED_DIR
+    if "GT_DIR" in globals():
+        gt_dir = GT_DIR
+
+    if prefixes is None:
+        prefixes = ["Drosophila", "MusMusculus", "Zebrafish"]
+
+    for prefix in prefixes:
+        conf_file = os.path.join(gt_dir, f"{prefix}_1_iitest_confidence.tif")
+        annot_file = os.path.join(gt_dir, f"{prefix}_1_iitest_annotations.tif")
+        conf_arr = tifffile.imread(conf_file)
+        gt_arr = tifffile.imread(annot_file)
+
+        seg_file = os.path.join(pred_dir, f"{prefix}_1_iitest_seg.tif")
+        seg_arr = tifffile.imread(seg_file)
+
+        # find largest overlap of ground truth mask with each segmentation instance
+        out_path = os.path.join(output_folder, f"{prefix}_matching_ref_gt.json")
+        if os.path.isfile(out_path) and not force_overwrite:
+            print(f"Skipping the creation of {out_path}. File already exists.")
+        else:
+            # exclude detections with more than 50% of pixels in undefined category
+            if 1 in np.unique(conf_arr)[1:]:
+                labels_undefined_mask = find_overlapping_masks(conf_arr, seg_arr, label_id_base=1)
+            else:
+                labels_undefined_mask = []
+                print("Array does not contain undefined mask")
+
+            find_matching_masks(gt_arr, seg_arr, out_path, labels_undefined_mask=labels_undefined_mask)
+
+        # find largest overlap of segmentation instance with each ground truth mask
+        out_path = os.path.join(output_folder, f"{prefix}_matching_gt_ref.json")
+        if os.path.isfile(out_path) and not force_overwrite:
+            print(f"Skipping the creation of {out_path}. File already exists.")
+        else:
+            find_matching_masks(seg_arr, gt_arr, out_path)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_folder", "-o", required=True)
+    parser.add_argument("--prefix", "-p", nargs="+", type=str, default=None)
+    parser.add_argument("--force", action="store_true", help="Forcefully overwrite output.")
+    args = parser.parse_args()
+
+    filter_true_positives(
+        args.output_folder,
+        args.prefix,
+        args.force,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/baselines/NIS3D_preprocessing.py

@@ -0,0 +1,139 @@
+import os
+
+import numpy as np
+import tifffile
+
+NIS3D_DIR = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/nucleus/NIS3D"
+TRAIN_DIR = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/nucleus/2025-07_NIS3D/train"
+VAL_DIR = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/nucleus/2025-07_NIS3D/val"
+TEST_DIR = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/nucleus/2025-07_NIS3D(test"
+
+# ---Training data---
+
+# clear: contains only 2,3,4 as seg ids
+train_dict_01 = {
+    "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/cross-image/train"),
+    "name": "Drosophila_2",
+    "conf_file": "ConfidenceScore.tif",
+    "gt_file": "GroundTruth.tif",
+    "output_dir": TRAIN_DIR,
+    "output_name": "Drosophila_2_annotations.tif",
+}
+
+# contains 1, 2, 3, 4
+train_dict_02 = {
+    "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/cross-image/train"),
+    "name": "Zebrafish_2",
+    "conf_file": "ConfidenceScore.tif",
+    "gt_file": "GroundTruth.tif",
+    "output_dir": TRAIN_DIR,
+    "output_name": "Zebrafish_2_annotations.tif",
+}
+
+# contains 1, 3, 4
+train_dict_03 = {
+   "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/cross-image/train"),
+   "name": "MusMusculus_2",
+   "conf_file": "scoreOfConfidence.tif",
+   "gt_file": "gt.tif",
+   "output_dir": TRAIN_DIR,
+   "output_name": "MusMusculus_2_annotations.tif",
+}
+
+# ---Validation data---
+
+val_dict_01 = {
+    "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/in-image/train"),
+    "name": "Drosophila_1",
+    "conf_file": "ConfidenceScore.tif",
+    "gt_file": "GroundTruth.tif",
+    "output_dir": VAL_DIR,
+    "output_name": "Drosophila_1_iitrain_annotations.tif",
+}
+
+val_dict_02 = {
+    "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/in-image/train"),
+    "name": "Zebrafish_1",
+    "conf_file": "ConfidenceScore.tif",
+    "gt_file": "GroundTruth.tif",
+    "output_dir": VAL_DIR,
+    "output_name": "Zebrafish_1_iitrain_annotations.tif",
+}
+
+val_dict_03 = {
+    "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/in-image/train"),
+    "name": "MusMusculus_1",
+    "conf_file": "ConfidenceScore.tif",
+    "gt_file": "GroundTruth.tif",
+    "output_dir": VAL_DIR,
+    "output_name": "MusMusculus_1_iitrain_annotations.tif",
+}
+
+# ---Test data---
+
+test_dict_01 = {
+    "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/in-image/test"),
+    "name": "Drosophila_1",
+    "conf_file": "ConfidenceScore.tif",
+    "gt_file": "GroundTruth.tif",
+    "output_dir": TEST_DIR,
+    "output_name": "Drosophila_1_iitest_annotations.tif",
+}
+
+test_dict_02 = {
+    "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/in-image/test"),
+    "name": "Zebrafish_1",
+    "conf_file": "ConfidenceScore.tif",
+    "gt_file": "GroundTruth.tif",
+    "output_dir": TEST_DIR,
+    "output_name": "Zebrafish_1_iitest_annotations.tif",
+}
+
+test_dict_03 = {
+    "data_dir": os.path.join(NIS3D_DIR, "suggestive_splitting/in-image/test"),
+    "name": "MusMusculus_1",
+    "conf_file": "ConfidenceScore.tif",
+    "gt_file": "GroundTruth.tif",
+    "output_dir": TEST_DIR,
+    "output_name": "MusMusculus_1_iitest_annotations.tif",
+}
+
+
+def filter_unmasked_data(conf_path, in_path, out_path):
+    conf = tifffile.imread(conf_path)
+    gt = tifffile.imread(in_path)
+    segmentation_ids = list(np.unique(conf)[1:])
+    if 1 in segmentation_ids:
+        instance_ids = list(np.unique(gt)[1:])
+        print(f"Number of instances before filtering: {len(instance_ids)}")
+        gt[conf == 1] = 0
+        instance_ids = list(np.unique(gt)[1:])
+        print(f"Number of instances after filtering: {len(instance_ids)}")
+        tifffile.imwrite(out_path, gt)
+    else:
+        instance_ids = list(np.unique(gt)[1:])
+        print(f"Number of instances: {len(instance_ids)}")
+        tifffile.imwrite(out_path, gt)
+
+
+def process_data_dicts(data_dicts):
+    for data_dict in data_dicts:
+        data_dir = data_dict["data_dir"]
+        dataset = os.path.join(data_dir, data_dict["name"])
+        conf_path = os.path.join(dataset, data_dict["conf_file"])
+        gt_path = os.path.join(dataset, data_dict["gt_file"])
+
+        out_dir = data_dict["output_dir"]
+        out_name = data_dict["output_name"]
+        out_path = os.path.join(out_dir, out_name)
+        filter_unmasked_data(conf_path, in_path=gt_path, out_path=out_path)
+
+
+def prepare_training_data():
+    """Prepare training data based on NIS3D data.
+
+    Cross-image data of half of the samples is used for training.
+    The other half of the samples is divided into validation data used for training and test data.
+    The in-image data is used for this, so that every remaining sample is split in half.
+    """
+    process_data_dicts([test_dict_01, test_dict_02, test_dict_03])

scripts/baselines/NIS3D_train.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+export MODEL_NAME="nucleus_NIS3D_supervised_2025-07-17"
+
+export IDIR=/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/nucleus/2025-07_NIS3D
+
+export SCRIPT_DIR=/user/schilling40/u15000/flamingo-tools/scripts/training
+
+python $SCRIPT_DIR/train_distance_unet.py -i $IDIR --name $MODEL_NAME
+