Fix issues in training CLI and add domain adaptation CLI

Author: constantinpape

scripts/cooper/revision/az_prediction.py

@@ -24,7 +24,7 @@ def run_prediction(model, name, split_folder, version, split_names, in_path):
 
     for fname in tqdm(file_names):
         if in_path:
-            input_path=os.path.join(in_path, name, fname)
+            input_path = os.path.join(in_path, name, fname)
         else:
             input_path = os.path.join(INPUT_ROOT, name, fname)
         print(f"segmenting {input_path}")
@@ -50,15 +50,14 @@ def run_prediction(model, name, split_folder, version, split_names, in_path):
                 print(f"{output_key_seg} already saved")
             else:
                 f.create_dataset(output_key_seg, data=seg, compression="lzf")
-                
 
 
 def get_model(version):
     assert version in (3, 4, 5, 6, 7)
     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")
-    elif version ==6:
+    elif version == 6:
         model_path = "/mnt/ceph-hdd/cold/nim00007/models/AZ/v6/"
     elif version == 7:
         model_path = "/mnt/lustre-emmy-hdd/usr/u12095/synapse_net/models/ConstantinAZ/checkpoints/v7/"
@@ -79,15 +78,15 @@ def main():
     args = parser.parse_args()
 
     if args.model_path:
-        model = load_model(model_path)
+        model = load_model(args.model_path)
     else:
         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, args.input)
-    
+
     print("Finished segmenting!")
 
 

scripts/cooper/revision/common.py

@@ -65,7 +65,7 @@ def get_split_folder(version):
     if version == 3:
         split_folder = "splits"
     elif version == 6:
-        split_folder= "/mnt/ceph-hdd/cold/nim00007/new_AZ_train_data/splits"
+        split_folder = "/mnt/ceph-hdd/cold/nim00007/new_AZ_train_data/splits"
     else:
         split_folder = "models_az_thin"
     return split_folder

setup.py

@@ -17,6 +17,7 @@
             "synapse_net.export_to_imod_points = synapse_net.tools.cli:imod_point_cli",
             "synapse_net.export_to_imod_objects = synapse_net.tools.cli:imod_object_cli",
             "synapse_net.run_supervised_training = synapse_net.training.supervised_training:main",
+            "synapse_net.run_domain_adaptation = synapse_net.training.domain_adaptation:main",
         ],
         "napari.manifest": [
             "synapse_net = synapse_net:napari.yaml",

synapse_net/training/domain_adaptation.py

@@ -1,12 +1,20 @@
 import os
+import tempfile
+from glob import glob
+from pathlib import Path
 from typing import Optional, Tuple
 
+import mrcfile
 import torch
 import torch_em
 import torch_em.self_training as self_training
+from elf.io import open_file
+from sklearn.model_selection import train_test_split
 
 from .semisupervised_training import get_unsupervised_loader
 from .supervised_training import get_2d_model, get_3d_model, get_supervised_loader, _determine_ndim
+from ..inference.inference import get_model_path, compute_scale_from_voxel_size
+from ..inference.util import _Scaler
 
 
 def mean_teacher_adaptation(
@@ -91,7 +99,7 @@ def mean_teacher_adaptation(
         if os.path.isdir(source_checkpoint):
             model = torch_em.util.load_model(source_checkpoint)
         else:
-            model = torch.load(source_checkpoint)
+            model = torch.load(source_checkpoint, weights_only=False)
         reinit_teacher = False
 
     optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
@@ -148,3 +156,109 @@ def mean_teacher_adaptation(
         sampler=sampler,
     )
     trainer.fit(n_iterations)
+
+
+# TODO patch shapes for other models
+PATCH_SHAPES = {
+    "vesicles_3d": [48, 256, 256],
+}
+"""@private
+"""
+
+
+def _get_paths(input_folder, pattern, resize_training_data, model_name, tmp_dir):
+    files = sorted(glob(os.path.join(input_folder, "**", pattern), recursive=True))
+    if len(files) == 0:
+        raise ValueError(f"Could not load any files from {input_folder} with pattern {pattern}")
+
+    val_fraction = 0.15
+
+    # Heuristic: if we have less then 4 files then we crop a part of the volumes for validation.
+    # And resave the volumes.
+    resave_val_crops = len(files) < 4
+
+    # We only resave the data if we resave val crops or resize the training data
+    resave_data = resave_val_crops or resize_training_data
+    if not resave_data:
+        train_paths, val_paths = train_test_split(files, test_size=val_fraction)
+        return train_paths, val_paths
+
+    train_paths, val_paths = [], []
+    for file_path in files:
+        file_name = os.path.basename(file_path)
+        data = open_file(file_path, mode="r")["data"][:]
+
+        if resize_training_data:
+            with mrcfile.open(file_path) as f:
+                voxel_size = f.voxel_size
+            voxel_size = {ax: vox_size / 10.0 for ax, vox_size in zip("xyz", voxel_size.item())}
+            scale = compute_scale_from_voxel_size(voxel_size, model_name)
+            scaler = _Scaler(scale, verbose=False)
+            data = scaler.sale_input(data)
+
+        if resave_val_crops:
+            n_slices = data.shape[0]
+            val_slice = int((1.0 - val_fraction) * n_slices)
+            train_data, val_data = data[:val_slice], data[val_slice:]
+
+            train_path = os.path.join(tmp_dir, Path(file_name).with_suffix(".h5")).replace(".h5", "_train.h5")
+            with open_file(train_path, mode="w") as f:
+                f.create_dataset("data", data=train_data, compression="lzf")
+            train_paths.append(train_path)
+
+            val_path = os.path.join(tmp_dir, Path(file_name).with_suffix(".h5")).replace(".h5", "_val.h5")
+            with open_file(val_path, mode="w") as f:
+                f.create_dataset("data", data=val_data, compression="lzf")
+            val_paths.append(val_path)
+
+        else:
+            output_path = os.path.join(tmp_dir, Path(file_name).with_suffix(".h5"))
+            with open_file(output_path, mode="w") as f:
+                f.create_dataset("data", data=data, compression="lzf")
+            train_paths.append(output_path)
+
+    if not resave_val_crops:
+        train_paths, val_paths = train_test_split(train_paths, test_size=val_fraction)
+
+    return train_paths, val_paths
+
+
+def _parse_patch_shape(patch_shape, model_name):
+    if patch_shape is None:
+        patch_shape = PATCH_SHAPES[model_name]
+    return patch_shape
+
+
+def main():
+    """@private
+    """
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description=""
+    )
+    parser.add_argument("--name", "-n", required=True)
+    parser.add_argument("--input", "-i", required=True)
+    parser.add_argument("--pattern", "-p", default="*.mrc")
+    parser.add_argument("--source_model", default="vesicles_3d")
+    parser.add_argument("--resize_training_data", action="store_true")
+    parser.add_argument("--n_iterations", type=int, default=int(1e4))
+    parser.add_argument("--patch_shape", nargs="+", type=int)
+    args = parser.parse_args()
+
+    source_checkpoint = get_model_path(args.source_model)
+    patch_shape = _parse_patch_shape(args.patch_shape, args.source_model)
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        unsupervised_train_paths, unsupervised_val_paths = _get_paths(
+            args.input, args.pattern, args.resize_training_data, args.source_model, tmp_dir
+        )
+
+        mean_teacher_adaptation(
+            name=args.name,
+            unsupervised_train_paths=unsupervised_train_paths,
+            unsupervised_val_paths=unsupervised_val_paths,
+            patch_shape=patch_shape,
+            source_checkpoint=source_checkpoint,
+            raw_key="data",
+            n_iterations=args.n_iterations,
+        )

synapse_net/training/supervised_training.py

@@ -307,9 +307,9 @@ def supervised_training(
 
 
 def _parse_input_folder(folder, pattern, key):
-    files = sorted(glob(os.path.join(folder, "**", pattern)))
+    files = sorted(glob(os.path.join(folder, "**", pattern), recursive=True))
     # Get all file extensions (general wild-cards may pick up files with multiple extensions).
-    extensions = [os.path.splitext(ff)[1] for ff in files]
+    extensions = list(set([os.path.splitext(ff)[1] for ff in files]))
 
     # If we have more than 1 file extension we just use the key that was passed,
     # as it is unclear how to derive a consistent key.
@@ -372,7 +372,7 @@ def main():
     parser.add_argument("-p", "--patch_shape", nargs=3, type=int, help="The patch shape for training.")
 
     # Folders with training data, containing raw/image data and labels.
-    parser.add_argument("--i", "--train_folder", required=True, help="The input folder with the training image data.")
+    parser.add_argument("-i", "--train_folder", required=True, help="The input folder with the training image data.")
     parser.add_argument("--image_file_pattern", default="*",
                         help="The pattern for selecting image files. For example, '*.mrc' to select all mrc files.")
     parser.add_argument("--raw_key",
@@ -394,6 +394,7 @@ def main():
     parser.add_argument("--n_samples_train", type=int, help="The number of samples per epoch for training. If not given will be derived from the data size.")  # noqa
     parser.add_argument("--n_samples_val", type=int, help="The number of samples per epoch for validation. If not given will be derived from the data size.")  # noqa
     parser.add_argument("--val_fraction", type=float, default=0.15, help="The fraction of the data to use for validation. This has no effect if 'val_folder' and 'val_label_folder' were passed.")  # noqa
+    parser.add_argument("--check", action="store_true", help="Visualize samples from the data loaders to ensure correct data instead of running training.")  # noqa
     args = parser.parse_args()
 
     train_image_paths, train_label_paths, val_image_paths, val_label_paths, raw_key, label_key =\