Domain adaptation for IHC segmentation of MLR215R (CR)

Author: schilling40

reproducibility/block_extraction/domain_adaptation_MLR215R_CR.json

@@ -0,0 +1,54 @@
+[
+	{
+		"cochlea": "M_LR_000215_R",
+		"image_channel": [
+			"CR"
+		],
+		"segmentation_channel": "IHC_v4b",
+		"type": "ihc",
+		"n_blocks": 6,
+		"halo_size": [
+			256,
+			256,
+			256
+		],
+		"component_list": [
+			7,
+			4,
+			1,
+			11
+		],
+		"crop_centers": [
+			[
+				727,
+				1300,
+				685
+			],
+			[
+				1123,
+				931,
+				870
+			],
+			[
+				812,
+				475,
+				1063
+			],
+			[
+				296,
+				645,
+				906
+			],
+			[
+				503,
+				918,
+				518
+			],
+			[
+				529,
+				1064,
+				408
+			]
+		]
+	}
+]

scripts/training/ihc_domain_adaptation_mouse_CR.py

@@ -0,0 +1,60 @@
+import os
+from glob import glob
+
+import torch
+from torch_em.util import load_model
+from flamingo_tools.training import mean_teacher_training
+
+
+def get_paths():
+    root = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/IHC"
+    folders = ["2025-08-IHC_domain_CR"]
+    train_paths = []
+    val_paths = []
+    for folder in folders:
+        train_paths.extend(glob(os.path.join(root, folder, "train", "*.tif")))
+        val_paths.extend(glob(os.path.join(root, folder, "val", "*.tif")))
+    return train_paths, val_paths
+
+
+def run_training(name):
+    patch_shape = (64, 128, 128)
+    batch_size = 8
+    source_checkpoint = os.path.join("/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet",
+                                     "trained_models/IHC",
+                                     "v4_cochlea_distance_unet_IHC_supervised_2025-07-14")
+
+    train_paths, val_paths = get_paths()
+    mean_teacher_training(
+        name=name,
+        unsupervised_train_paths=train_paths,
+        unsupervised_val_paths=val_paths,
+        patch_shape=patch_shape,
+        source_checkpoint=source_checkpoint,
+        batch_size=batch_size,
+        n_iterations=int(2.5e4),
+        n_samples_train=1000,
+        n_samples_val=80,
+    )
+
+
+def export_model(name, export_path):
+    model = load_model(os.path.join("checkpoints", name), state_key="teacher")
+    torch.save(model, export_path)
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--export_path")
+    args = parser.parse_args()
+    name = "ihc-adapted-model_M215R_CR"
+    if args.export_path is None:
+        run_training(name)
+    else:
+        export_model(name, args.export_path)
+
+
+if __name__ == "__main__":
+    main()