Clean up AZ scripts and summarize current state

Author: constantinpape

scripts/cooper/revision/README.md

@@ -0,0 +1,19 @@
+# Improving the AZ model
+
+Scripts for improving the AZ annotations, training the AZ model, and evaluating it.
+
+The most important scripts are:
+- For improving and updating the AZ annotations:
+    - `prediction.py`: Run prediction of vesicle and boundary model.
+    - `thin_az_gt.py`: Thin the AZ annotations, so that it aligns only with the presynaptic membrane. This is done by intersecting the annotations with the presynaptic compartment, using predictions from the network used for compartment segmentation.
+    - `assort_new_az_data.py`: Create a new version of the annotation, renaming the dataset, and creating a cropped version of the endbulb of held data.
+    - `merge_az.py`: Merge AZ annotations with predictions from model v4, in order to remove some artifacts that resulted from AZ thinning.
+- For evaluating the AZ predictions: 
+    - `az_prediction.py`: Run prediction with the AZ model.
+    - `run_az_evaluation.py`: Evaluate the predictions of an AZ model.
+    - `evaluate_result.py`: Summarize the evaluation results.
+- And for training: `train_az_gt.py`. So far, I have trained:
+    - v3: Trained on the initial annotations.
+    - v4: Trained on the thinned annotations.
+    - v5: Trained on the thinned annotations with an additional distance loss (did not help).
+    - v6: Trained on the merged annotations.

scripts/cooper/revision/az_prediction.py

@@ -0,0 +1,61 @@
+import argparse
+import os
+
+import h5py
+from synapse_net.inference.active_zone import segment_active_zone
+from torch_em.util import load_model
+from tqdm import tqdm
+
+from common import get_file_names, get_split_folder, ALL_NAMES, INPUT_ROOT, OUTPUT_ROOT
+
+
+def run_prediction(model, name, split_folder, version, split_names):
+    file_names = get_file_names(name, split_folder, split_names=split_names)
+
+    output_folder = os.path.join(OUTPUT_ROOT, name)
+    os.makedirs(output_folder, exist_ok=True)
+    output_key = f"predictions/az/v{version}"
+
+    for fname in tqdm(file_names):
+        output_path = os.path.join(output_folder, fname)
+
+        if os.path.exists(output_path):
+            with h5py.File(output_path, "r") as f:
+                if output_key in f:
+                    continue
+
+        input_path = os.path.join(INPUT_ROOT, name, fname)
+        with h5py.File(input_path, "r") as f:
+            raw = f["raw"][:]
+
+        _, pred = segment_active_zone(raw, model=model, verbose=False, return_predictions=True)
+        with h5py.File(output_path, "a") as f:
+            f.create_dataset(output_key, data=pred, compression="lzf")
+
+
+def get_model(version):
+    assert version in (3, 4, 5)
+    split_folder = get_split_folder(version)
+    if version == 3:
+        model_path = os.path.join(split_folder, "checkpoints", "3D-AZ-model-TEM_STEM_ChemFix_wichmann-v3")
+    else:
+        model_path = os.path.join(split_folder, "checkpoints", f"v{version}")
+    model = load_model(model_path)
+    return model
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--version", "-v", type=int)
+    parser.add_argument("--names", nargs="+", default=ALL_NAMES)
+    parser.add_argument("--splits", nargs="+", default=["test"])
+    args = parser.parse_args()
+
+    model = get_model(args.version)
+    split_folder = get_split_folder(args.version)
+    for name in args.names:
+        run_prediction(model, name, split_folder, args.version, args.splits)
+
+
+if __name__ == "__main__":
+    main()

scripts/cooper/revision/check_prediction.py

@@ -0,0 +1,44 @@
+import argparse
+import os
+
+import h5py
+import napari
+from common import ALL_NAMES, get_file_names, get_split_folder, get_paths
+
+
+def check_predictions(name, split, version):
+    split_folder = get_split_folder(version)
+    file_names = get_file_names(name, split_folder, split_names=[split])
+    seg_paths, gt_paths = get_paths(name, file_names)
+
+    for seg_path, gt_path in zip(seg_paths, gt_paths):
+
+        with h5py.File(gt_path, "r") as f:
+            raw = f["raw"][:]
+            gt = f["labels/az"][:] if version == 3 else f["labels/az_thin"][:]
+
+        with h5py.File(seg_path) as f:
+            seg_key = f"predictions/az/v{version}"
+            pred = f[seg_key][:]
+
+        v = napari.Viewer()
+        v.add_image(raw)
+        v.add_image(pred, blending="additive")
+        v.add_labels(gt)
+        v.title = f"{name}/{os.path.basename(seg_path)}"
+        napari.run()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--version", "-v", type=int, required=True)
+    parser.add_argument("--split", default="test")
+    parser.add_argument("--names", nargs="+", default=ALL_NAMES)
+    args = parser.parse_args()
+
+    for name in args.names:
+        check_predictions(name, args.split, args.version)
+
+
+if __name__ == "__main__":
+    main()

scripts/cooper/revision/common.py

@@ -0,0 +1,69 @@
+import json
+import os
+
+
+# The root folder which contains the new AZ training data.
+INPUT_ROOT = "/mnt/ceph-hdd/cold_store/projects/nim00007/new_AZ_train_data"
+# The output folder for AZ predictions.
+OUTPUT_ROOT = "/mnt/ceph-hdd/cold_store/projects/nim00007/AZ_predictions_new"
+
+# The names of all datasets for which to run prediction / evaluation.
+# This excludes 'endbulb_of_held_cropped', which is a duplicate of 'endbulb_of_held',
+# which we don't evaluate on because of this.
+ALL_NAMES = [
+    "chemical_fixation", "endbulb_of_held", "stem", "stem_cropped", "tem"
+]
+
+# The translation of new dataset names to old dataset names.
+NAME_TRANSLATION = {
+    "chemical_fixation": ["12_chemical_fix_cryopreparation_minusSVseg_corrected"],
+    "endbulb_of_held": ["wichmann_withAZ_rescaled_tomograms"],
+    "stem": ["04_hoi_stem_examples_fidi_and_sarah_corrected_rescaled_tomograms"],
+    "stem_cropped": ["04_hoi_stem_examples_minusSVseg_cropped_corrected_rescaled_tomograms",
+                     "06_hoi_wt_stem750_fm_minusSVseg_cropped_corrected_rescaled_tomograms"],
+    "tem": ["01data_withoutInvertedFiles_minusSVseg_corrected"],
+}
+
+
+# Get the paths to the files with raw data / ground-truth and the segmentation.
+def get_paths(name, file_names, skip_seg=False):
+    seg_paths, gt_paths = [], []
+    for fname in file_names:
+        if not skip_seg:
+            seg_path = os.path.join(OUTPUT_ROOT, name, fname)
+            assert os.path.exists(seg_path), seg_path
+            seg_paths.append(seg_path)
+
+        gt_path = os.path.join(INPUT_ROOT, name, fname)
+        assert os.path.exists(gt_path), gt_path
+        gt_paths.append(gt_path)
+
+    return seg_paths, gt_paths
+
+
+def get_file_names(name, split_folder, split_names):
+    split_path = os.path.join(split_folder, f"split-{name}.json")
+    if os.path.exists(split_path):
+        with open(split_path) as f:
+            splits = json.load(f)
+        file_names = [fname for split in split_names for fname in splits[split]]
+
+    else:
+        old_names = NAME_TRANSLATION[name]
+        file_names = []
+        for old_name in old_names:
+            split_path = os.path.join(split_folder, f"split-{old_name}.json")
+            with open(split_path) as f:
+                splits = json.load(f)
+            this_file_names = [fname for split in split_names for fname in splits[split]]
+            file_names.extend(this_file_names)
+    return file_names
+
+
+def get_split_folder(version):
+    assert version in (3, 4, 5)
+    if version == 3:
+        split_folder = "splits"
+    else:
+        split_folder = "models_az_thin"
+    return split_folder

scripts/cooper/revision/eval_AZ.sh

@@ -6,6 +6,7 @@
 args = parser.parse_args()
 
 results = pd.read_excel(args.result_path)
+print(results)
 
 
 def summarize_results(res):
@@ -20,29 +21,29 @@ def summarize_results(res):
 
 
 # # Compute the results for Chemical Fixation.
-results_chem_fix = results[results.dataset.str.startswith("12")]
+results_chem_fix = results[results.dataset == "chemical_fixation"]
 if results_chem_fix.size > 0:
     print("Chemical Fixation Results:")
     summarize_results(results_chem_fix)
 #
 # # Compute the results for STEM (=04).
-results_stem = results[results.dataset.str.startswith(("04", "06"))]
+results_stem = results[results.dataset.str.startswith("stem")]
 if results_stem.size > 0:
     print()
     print("STEM Results:")
     summarize_results(results_stem)
 #
 # # Compute the results for TEM (=01).
-results_tem = results[results.dataset.str.startswith("01")]
+results_tem = results[results.dataset == "tem"]
 if results_tem.size > 0:
     print()
     print("TEM Results:")
     summarize_results(results_tem)
 
 #
-# Compute the results for Wichmann.
-results_wichmann = results[results.dataset.str.startswith("wichmann")]
+# Compute the results for Wichmann / endbulb of held.
+results_wichmann = results[results.dataset.str.startswith("endbulb")]
 if results_wichmann.size > 0:
     print()
-    print("Wichmann Results:")
+    print("Endbulb of Held Results:")
     summarize_results(results_wichmann)

scripts/cooper/revision/evaluate_result.py

@@ -0,0 +1,17 @@
+import os
+from glob import glob
+import h5py
+from tqdm import tqdm
+
+
+INPUT_ROOT = "/mnt/ceph-hdd/cold_store/projects/nim00007/new_AZ_train_data"
+
+files = glob(os.path.join(INPUT_ROOT, "**/*.h5"), recursive=True)
+
+key = "labels/az_merged"
+for ff in tqdm(files):
+    with h5py.File(ff, "a") as f:
+        az = f[key][:]
+        az = az.squeeze()
+        del f[key]
+        f.create_dataset(key, data=az, compression="lzf")