Update structure segmentation training

Author: constantinpape

scripts/inner_ear/training/train_structure_segmentation.py

@@ -10,7 +10,22 @@
 from tqdm import tqdm
 
 ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/moser/inner_ear_data"
+ROOT_OTHER_TOMOS = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/moser/other_tomograms/"
+
 LABEL_KEY = "labels/inner_ear_structures"
+OTHER_NAMES = ["vesicle_pools", "tether", "rat"]
+
+
+def get_other_paths(name):
+    assert name in OTHER_NAMES, f"Invalid name {name}"
+    if name == "vesicle_pools":
+        folder = "01_vesicle_pools"
+    elif name == "tether":
+        folder = "02_tether"
+    else:
+        folder = "03_ratten_tomos"
+    paths = sorted(glob(os.path.join(ROOT, folder, "*.h5")))
+    return paths
 
 
 def get_train_val_test_split(root):
@@ -34,8 +49,7 @@ def get_train_val_test_split(root):
     return train_tomos, val_tomos, test_tomos
 
 
-def preprocess_labels(tomograms):
-    structure_keys = ("ribbon", "PD", "membrane")
+def preprocess_labels(tomograms, structure_keys=("ribbon", "PD", "membrane")):
     nc = len(structure_keys)
 
     for tomo in tqdm(tomograms, desc="Preprocess labels"):
@@ -60,11 +74,11 @@ def noop(x):
     return x
 
 
-def train_inner_ear_structures(train_tomograms, val_tomograms):
+def train_inner_ear_structures(train_tomograms, val_tomograms, name):
     patch_shape = (64, 512, 512)
     sampler = MinForegroundSampler(min_fraction=0.05, p_reject=0.95)
     supervised_training(
-        name="inner_ear_structure_model",
+        name=name,
         train_paths=train_tomograms, val_paths=val_tomograms,
         label_key=LABEL_KEY, patch_shape=patch_shape, save_root=".",
         sampler=sampler, label_transform=noop, out_channels=3,
@@ -73,11 +87,27 @@ def train_inner_ear_structures(train_tomograms, val_tomograms):
     )
 
 
-def main():
+def training_v1():
+    train_tomograms, val_tomograms, _ = get_train_val_test_split(ROOT)
+    preprocess_labels(train_tomograms)
+    preprocess_labels(val_tomograms)
+    train_inner_ear_structures(train_tomograms, val_tomograms, name="inner_ear_structure_model")
+
+
+def training_v2():
     train_tomograms, val_tomograms, _ = get_train_val_test_split(ROOT)
     preprocess_labels(train_tomograms)
+    for name in OTHER_NAMES:
+        other_tomograms = get_other_paths(name)
+        preprocess_labels(other_tomograms, structure_keys=("ribbons", "presynapse", "membrane"))
+        train_tomograms.extend(other_tomograms)
     preprocess_labels(val_tomograms)
-    train_inner_ear_structures(train_tomograms, val_tomograms)
+    train_inner_ear_structures(train_tomograms, val_tomograms, name="inner_ear_structure_model_v2")
+
+
+def main():
+    # training_v1()
+    training_v2()
 
 
 if __name__ == "__main__":