Update synapse validation

Author: constantinpape

flamingo_tools/segmentation/synapse_detection.py

@@ -25,7 +25,7 @@ def map_and_filter_detections(
     Args:
         segmentation: The IHC segmentation.
         detections: The synapse marker detections.
-        max_distance: The maximal distance for a valid match of synapse markers to IHCs.
+        max_distance: The maximal distance in micrometer for a valid match of synapse markers to IHCs.
         resolution: The resolution / voxel size of the data in micrometer.
         n_threads: The number of threads for parallelizing the mapping of detections to objects.
         verbose: Whether to print the progress of the mapping procedure.

scripts/synapse_marker_detection/extract_training_data.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import h5py
+import imageio.v3 as imageio
 import napari
 import numpy as np
 import pandas as pd
@@ -19,34 +20,61 @@ def get_voxel_size(imaris_file):
     return vsize
 
 
-def extract_training_data(imaris_file, output_folder, crop=True, scale=True):
+def get_transformation(imaris_file):
+    with h5py.File(imaris_file) as f:
+        info = f["DataSetInfo"]["Image"].attrs
+        ext_min = np.array([float(b"".join(info[f"ExtMin{i}"]).decode()) for i in range(3)])
+        ext_max = np.array([float(b"".join(info[f"ExtMax{i}"]).decode()) for i in range(3)])
+        size = [int(b"".join(info[dim]).decode()) for dim in ["X", "Y", "Z"]]
+        spacing = (ext_max - ext_min) / size                              # µm / voxel
+
+    # build 4×4 affine: world → index
+    T = np.eye(4)
+    T[:3, :3] = np.diag(1/spacing)            # scale
+    T[:3, 3] = -ext_min/spacing              # translate
+
+    return T
+
+
+def extract_training_data(imaris_file, output_folder, tif_file=None, crop=True):
     point_key = "/Scene/Content/Points0/CoordsXYZR"
     with h5py.File(imaris_file, "r") as f:
         if point_key not in f:
             print("Skipping", imaris_file, "due to missing annotations")
             return
-        data = f["/DataSet/ResolutionLevel 0/TimePoint 0/Channel 0/Data"][:]
         points = f[point_key][:]
         points = points[:, :-1]
-        points = points[:, ::-1]
 
-    # TODO crop the data to the original shape.
-    # Can we just crop the zero-padding ?!
+        g = f["/DataSet/ResolutionLevel 0/TimePoint 0"]
+        # The first channel is ctbp2 / the synapse marker channel.
+        data = g["Channel 0/Data"][:]
+        # The second channel is vglut / the ihc channel.
+        if "Channel 1" in g:
+            ihc_data = g["Channel 1/Data"][:]
+        else:
+            ihc_data = None
+
+    T = get_transformation(imaris_file)
+    points = (T @ np.c_[points, np.ones(len(points))].T).T[:, :3]
+    points = points[:, ::-1]
+
     if crop:
         crop_box = np.where(data != 0)
         crop_box = tuple(slice(0, int(cb.max() + 1)) for cb in crop_box)
         data = data[crop_box]
 
-    # Scale the points to match the image dimensions.
-    voxel_size = get_voxel_size(imaris_file)
-    if scale:
-        points /= voxel_size[None]
-
-    print(data.shape, voxel_size)
+    if tif_file is None:
+        original_data = None
+    else:
+        original_data = imageio.imread(tif_file)
 
     if output_folder is None:
         v = napari.Viewer()
         v.add_image(data)
+        if ihc_data is not None:
+            v.add_image(ihc_data)
+        if original_data is not None:
+            v.add_image(original_data, visible=False)
         v.add_points(points)
         v.title = os.path.basename(imaris_file)
         napari.run()
@@ -66,6 +94,8 @@ def extract_training_data(imaris_file, output_folder, crop=True, scale=True):
 
         f = zarr.open(image_file, "a")
         f.create_dataset("raw", data=data)
+        if ihc_data is not None:
+            f.create_dataset("raw_ihc", data=ihc_data)
 
 
 # Files that look good for training:
@@ -82,6 +112,21 @@ def process_training_data_v1():
         extract_training_data(ff, output_folder="./training_data")
 
 
+def _match_tif(imaris_file):
+    folder = os.path.split(imaris_file)[0]
+
+    fname = os.path.basename(imaris_file)
+    parts = fname.split("_")
+    cochlea = parts[0].upper()
+    region = parts[1]
+
+    tif_name = f"{cochlea}_{region}_CTBP2.tif"
+    tif_path = os.path.join(folder, tif_name)
+    assert os.path.exists(tif_path), tif_path
+
+    return tif_path
+
+
 def process_training_data_v2(visualize=True):
     input_root = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/ImageCropsIHC_synapses"
 
@@ -110,16 +155,46 @@ def process_training_data_v2(visualize=True):
 
         imaris_files = sorted(glob(os.path.join(input_root, folder, "*.ims")))
         for imaris_file in imaris_files:
-            fname = os.path.basename(imaris_file)
-            if fname not in valid_files:
+            if os.path.basename(imaris_file) not in valid_files:
+                continue
+            extract_training_data(imaris_file, output_folder, tif_file=None, crop=True, scale=True)
+
+
+# We have fixed the imaris data extraction problem and can use all the crops!
+def process_training_data_v3(visualize=True):
+    input_root = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/ImageCropsIHC_synapses"
+
+    train_output = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/training_data/v3"  # noqa
+    test_output = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/test_data/v3"  # noqa
+
+    train_folders = ["synapse_stains", "M78L_IHC-synapse_crops", "M226R_IHC-synapsecrops"]
+    test_folders = ["M226L_IHC-synapse_crops"]
+
+    exclude_names = ["220824_Ex3IL_rbCAST1635_mCtBP2580_chCR488_cell1_CtBP2spots.ims"]
+
+    for folder in train_folders + test_folders:
+
+        if visualize:
+            output_folder = None
+        elif folder in train_folders:
+            output_folder = train_output
+            os.makedirs(output_folder, exist_ok=True)
+        else:
+            output_folder = test_output
+            os.makedirs(output_folder, exist_ok=True)
+
+        imaris_files = sorted(glob(os.path.join(input_root, folder, "*.ims")))
+        for imaris_file in imaris_files:
+            if os.path.basename(imaris_file) in exclude_names:
+                print("Skipping", imaris_file)
                 continue
-            print(fname)
-            extract_training_data(imaris_file, output_folder, crop=True, scale=True)
+            extract_training_data(imaris_file, output_folder, tif_file=None, crop=True)
 
 
 def main():
     # process_training_data_v1()
-    process_training_data_v2(visualize=False)
+    # process_training_data_v2(visualize=True)
+    process_training_data_v3(visualize=False)
 
 
 if __name__ == "__main__":

scripts/synapse_marker_detection/train_synapse_detection.py

@@ -1,49 +1,42 @@
 import os
 import sys
+from glob import glob
 
+from sklearn.model_selection import train_test_split
 from detection_dataset import DetectionDataset, MinPointSampler
 
 sys.path.append("/home/pape/Work/my_projects/czii-protein-challenge")
 sys.path.append("/user/pape41/u12086/Work/my_projects/czii-protein-challenge")
 
 from utils.training.training import supervised_training  # noqa
 
-ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/training_data/v2"  # noqa
+ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/training_data/v3"  # noqa
 TRAIN_ROOT = os.path.join(ROOT, "images")
 LABEL_ROOT = os.path.join(ROOT, "labels")
 
 
 def get_paths(split):
-    file_names = [
-        "4.1L_apex_IHCribboncount_Z",
-        "4.1L_base_IHCribbons_Z",
-        "4.1L_mid_IHCribboncount_Z",
-        "4.2R_apex_IHCribboncount_Z",
-        "4.2R_apex_IHCribboncount_Z",
-        "6.2R_apex_IHCribboncount_Z",
-        "m78l_apexp2718_cr-ctbp2",
-        "6.2R_base_IHCribbons_Z",
-    ]
-    image_paths = [os.path.join(TRAIN_ROOT, f"{fname}.zarr") for fname in file_names]
-    label_paths = [os.path.join(LABEL_ROOT, f"{fname}.csv") for fname in file_names]
+    image_paths = sorted(glob(os.path.join(TRAIN_ROOT, "*.zarr")))
+    label_paths = sorted(glob(os.path.join(LABEL_ROOT, "*.csv")))
+    assert len(image_paths) == len(label_paths)
+
+    train_images, val_images, train_labels, val_labels = train_test_split(
+        image_paths, label_paths, test_size=2, random_state=42
+    )
 
     if split == "train":
-        image_paths = image_paths[:-1]
-        label_paths = label_paths[:-1]
+        image_paths = train_images
+        label_paths = train_labels
     else:
-        image_paths = image_paths[-1:]
-        label_paths = label_paths[-1:]
-
-    for path in image_paths:
-        assert os.path.exists(path), path
+        image_paths = val_images
+        label_paths = val_labels
 
     return image_paths, label_paths
 
 
-# TODO maybe add a sampler for the label data
 def train():
 
-    model_name = "synapse_detection_v2"
+    model_name = "synapse_detection_v3"
 
     train_paths, train_label_paths = get_paths("train")
     val_paths, val_label_paths = get_paths("val")
@@ -56,7 +49,7 @@ def train():
 
     patch_shape = [40, 112, 112]
     batch_size = 32
-    check = True
+    check = False
 
     supervised_training(
         name=model_name,
@@ -81,7 +74,7 @@ def train():
         dataset_class=DetectionDataset,
         n_samples_train=3200,
         n_samples_val=160,
-        sampler=MinPointSampler(min_points=1, p_reject=0.6),
+        sampler=MinPointSampler(min_points=1, p_reject=0.8),
     )
 
 

scripts/validation/synapses/prediction.py

@@ -10,15 +10,16 @@
 from elf.parallel.local_maxima import find_local_maxima
 from flamingo_tools.segmentation.unet_prediction import prediction_impl, run_unet_prediction
 
-INPUT_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/test_data/v2/images"  # noqa
+INPUT_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/test_data/v3/images"  # noqa
+GT_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/test_data/v3/labels"
 OUTPUT_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/AnnotatedImageCrops/SynapseValidation"
 
 sys.path.append("/user/pape41/u12086/Work/my_projects/czii-protein-challenge")
 sys.path.append("../../synapse_marker_detection")
 
 
 def pred_synapse_impl(input_path, output_folder):
-    model_path = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/synapse_marker_detection/checkpoints/synapse_detection_v2"  # noqa
+    model_path = "/mnt/vast-nhr/home/pape41/u12086/Work/my_projects/flamingo-tools/scripts/synapse_marker_detection/checkpoints/synapse_detection_v3"  # noqa
     input_key = "raw"
 
     block_shape = (32, 128, 128)
@@ -48,7 +49,7 @@ def pred_synapse_impl(input_path, output_folder):
 
 
 def predict_synapses():
-    files = glob(os.path.join(INPUT_ROOT, "*.zarr"))
+    files = sorted(glob(os.path.join(INPUT_ROOT, "*.zarr")))
     for ff in files:
         print("Segmenting", ff)
         output_folder = os.path.join(OUTPUT_ROOT, Path(ff).stem)
@@ -59,34 +60,95 @@ def pred_ihc_impl(input_path, output_folder):
     model_path = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/trained_models/IHC/v2_cochlea_distance_unet_IHC_supervised_2025-05-21"  # noqa
 
     run_unet_prediction(
-        input_path, input_key=None, output_folder=output_folder, model_path=model_path, min_size=1000,
+        input_path, input_key="raw_ihc", output_folder=output_folder, model_path=model_path, min_size=1000,
         seg_class="ihc", center_distance_threshold=0.5, boundary_distance_threshold=0.5,
     )
 
 
 def predict_ihcs():
-    files = [
-        "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/ImageCropsIHC_synapses/M226R_IHC-synapsecrops/M226R_base_p800_Vglut3.tif",  # noqa
-        "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/ImageCropsIHC_synapses/M226R_IHC-synapsecrops/M226R_apex_p1268_Vglut3.tif",  # noqa
-    ]
+    files = sorted(glob(os.path.join(INPUT_ROOT, "*.zarr")))
     for ff in files:
         print("Segmenting", ff)
-        output_folder = os.path.join(OUTPUT_ROOT, Path(ff).stem)
+        output_folder = os.path.join(OUTPUT_ROOT, f"{Path(ff).stem}_ihc")
         pred_ihc_impl(ff, output_folder)
 
 
-# TODO also filter GT
+def _filter_synapse_impl(detections, ihc_file, output_path):
+    from flamingo_tools.segmentation.synapse_detection import map_and_filter_detections
+
+    with open_file(ihc_file, mode="r") as f:
+        if "segmentation_filtered" in f:
+            print("Using filtered segmentation!")
+            segmentation = open_file(ihc_file)["segmentation_filtered"][:]
+        else:
+            segmentation = open_file(ihc_file)["segmentation"][:]
+
+    max_distance = 5  # 5 micrometer
+    filtered_detections = map_and_filter_detections(segmentation, detections, max_distance=max_distance)
+    filtered_detections.to_csv(output_path, index=False, sep="\t")
+
+
 def filter_synapses():
-    pass
+    input_files = sorted(glob(os.path.join(INPUT_ROOT, "*.zarr")))
+    for ff in input_files:
+        ihc = os.path.join(OUTPUT_ROOT, f"{Path(ff).stem}_ihc", "segmentation.zarr")
+        output_folder = os.path.join(OUTPUT_ROOT, Path(ff).stem)
+        synapses = os.path.join(output_folder, "synapse_detection.tsv")
+        synapses = pd.read_csv(synapses, sep="\t")
+        output_path = os.path.join(output_folder, "filtered_synapse_detection.tsv")
+        _filter_synapse_impl(synapses, ihc, output_path)
+
+
+def filter_gt():
+    input_files = sorted(glob(os.path.join(INPUT_ROOT, "*.zarr")))
+    gt_files = sorted(glob(os.path.join(GT_ROOT, "*.csv")))
+    for ff, gt in zip(input_files, gt_files):
+        ihc = os.path.join(OUTPUT_ROOT, f"{Path(ff).stem}_ihc", "segmentation.zarr")
+        output_folder, fname = os.path.split(gt)
+        output_path = os.path.join(output_folder, fname.replace(".csv", "_filtered.tsv"))
+
+        gt = pd.read_csv(gt)
+        gt = gt.rename(columns={"axis-0": "z", "axis-1": "y", "axis-2": "x"})
+        gt.insert(0, "spot_id", np.arange(1, len(gt) + 1))
+
+        _filter_synapse_impl(gt, ihc, output_path)
+
+
+def _check_prediction(input_file, ihc_file, detection_file):
+    import napari
+
+    synapses = pd.read_csv(detection_file, sep="\t")[["z", "y", "x"]].values
+
+    vglut = open_file(input_file)["raw_ihc"][:]
+    ctbp2 = open_file(input_file)["raw"][:]
+    ihcs = open_file(ihc_file)["segmentation"][:]
+
+    v = napari.Viewer()
+    v.add_image(vglut)
+    v.add_image(ctbp2)
+    v.add_labels(ihcs)
+    v.add_points(synapses)
+    napari.run()
 
 
 def check_predictions():
-    pass
+    input_files = sorted(glob(os.path.join(INPUT_ROOT, "*.zarr")))
+    for ff in input_files:
+        ihc = os.path.join(OUTPUT_ROOT, f"{Path(ff).stem}_ihc", "segmentation.zarr")
+        synapses = os.path.join(OUTPUT_ROOT, Path(ff).stem, "filtered_synapse_detection.tsv")
+        _check_prediction(ff, ihc, synapses)
 
 
-def main():
-    # predict_synapses()
+def process_everything():
+    predict_synapses()
     predict_ihcs()
+    filter_synapses()
+    filter_gt()
+
+
+def main():
+    process_everything()
+    # check_predictions()
 
 
 if __name__ == "__main__":