#!/usr/bin/env python3 """ train_siamese.py — Train a siamese CNN on Task06_Lung to determine if two CT slices come from the same patient. Usage: conda activate fundus_imaging python scripts/train_siamese.py python scripts/train_siamese.py --backbone resnet34 --epochs 15 """ import os import sys import json import time import argparse import numpy as np os.environ["OMP_NUM_THREADS"] = "1" os.environ["OPENBLAS_NUM_THREADS"] = "1" os.environ["MKL_NUM_THREADS"] = "1" import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image from sklearn.metrics import accuracy_score, roc_auc_score from scipy.sparse.csgraph import connected_components from scipy.sparse import csr_matrix from collections import defaultdict sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from classes import NiftiSliceDataset, SiameseCNN # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) MODELS_DIR = os.path.join(ROOT, "models") RESULTS_DIR = os.path.join(ROOT, "results") DEFAULT_CACHE_DIR = os.path.join(ROOT, "features", "task06_pngs") BACKBONES = { "resnet18": (512, 224), "resnet34": (512, 224), "efficientnet_b0": (1280, 240), } SEED = 42 BATCH_SIZE = 64 EPOCHS = 15 LR = 1e-4 MIN_Z_GAP = 20 HARD_NEG_FRAC = 0.5 TEST_PATIENTS = 8 N_PAIRS = 20000 # --------------------------------------------------------------------------- # Pair dataset # --------------------------------------------------------------------------- class PairDataset(Dataset): """Yield (img_A, img_B, label) pairs with hard negative mining.""" def __init__(self, paths, patient_ids, z_indices, n_pairs, transform=None, min_z_gap=MIN_Z_GAP, hard_neg_frac=HARD_NEG_FRAC, seed=SEED): rng = np.random.default_rng(seed) unique_patients = np.unique(patient_ids) z_to_samples = defaultdict(list) for idx in range(len(paths)): z_to_samples[z_indices[idx]].append((idx, patient_ids[idx])) n_pos = n_pairs // 2 n_neg = n_pairs - n_pos n_hard = int(n_neg * hard_neg_frac) pairs, labels = [], [] pos_z_pairs = [] # Positive pairs: same patient, z-distance >= min_z_gap for _ in range(n_pos): for __ in range(100): pid = rng.choice(unique_patients) idx = np.where(patient_ids == pid)[0] if len(idx) < 2: continue i, j = rng.choice(idx, size=2, replace=False) if abs(int(z_indices[i]) - int(z_indices[j])) >= min_z_gap: pairs.append((i, j)) labels.append(1) pos_z_pairs.append((z_indices[i], z_indices[j])) break while len(labels) < n_pos: pid = rng.choice(unique_patients) idx = np.where(patient_ids == pid)[0] if len(idx) < 2: continue i, j = rng.choice(idx, size=2, replace=False) pairs.append((i, j)) labels.append(1) pos_z_pairs.append((z_indices[i], z_indices[j])) # Hard negatives: different patients, matched z-positions for _ in range(n_hard): z_i, z_j = pos_z_pairs[rng.integers(0, len(pos_z_pairs))] for __ in range(100): s_i = z_to_samples.get(z_i, []) s_j = z_to_samples.get(z_j, []) if len(s_i) < 1 or len(s_j) < 1: break si = s_i[rng.integers(0, len(s_i))] sj = s_j[rng.integers(0, len(s_j))] if si[1] != sj[1]: pairs.append((si[0], sj[0])) labels.append(0) break # Easy negatives while len(labels) < n_pairs: p1, p2 = rng.choice(unique_patients, size=2, replace=False) i = rng.choice(np.where(patient_ids == p1)[0]) j = rng.choice(np.where(patient_ids == p2)[0]) pairs.append((i, j)) labels.append(0) order = rng.permutation(len(labels)) self.pairs = [(pairs[o][0], pairs[o][1]) for o in order] self.labels = [labels[o] for o in order] self.paths = paths self.transform = transform def __len__(self): return len(self.pairs) def __getitem__(self, idx): i, j = self.pairs[idx] img_a = Image.open(self.paths[i]).convert("RGB") img_b = Image.open(self.paths[j]).convert("RGB") if self.transform: img_a = self.transform(img_a) img_b = self.transform(img_b) return img_a, img_b, torch.tensor(self.labels[idx], dtype=torch.float32) # --------------------------------------------------------------------------- # Training # --------------------------------------------------------------------------- def train_epoch(model, loader, optimizer, criterion, device): model.train() total_loss, correct, n = 0.0, 0, 0 for a, b, y in loader: a, b, y = a.to(device), b.to(device), y.to(device) optimizer.zero_grad() loss = criterion(model(a, b), y) loss.backward() optimizer.step() preds = (torch.sigmoid(model(a, b)) > 0.5).float() total_loss += loss.item() * len(y) correct += (preds == y).sum().item() n += len(y) return total_loss / n, correct / n @torch.no_grad() def eval_epoch(model, loader, criterion, device): model.eval() total_loss, n = 0.0, 0 all_preds, all_labels = [], [] for a, b, y in loader: a, b, y = a.to(device), b.to(device), y.to(device) logits = model(a, b) total_loss += criterion(logits, y).item() * len(y) all_preds.extend(torch.sigmoid(logits).cpu().numpy()) all_labels.extend(y.cpu().numpy()) n += len(y) all_labels = np.array(all_labels, dtype=int) all_preds = np.array(all_preds, dtype=float) acc = accuracy_score(all_labels, all_preds > 0.5) try: auc = roc_auc_score(all_labels, all_preds) except ValueError: auc = 0.5 return total_loss / n, acc, auc # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _export_subset(slices, patient_ids, filenames, cache_dir): """Export a subset of slices as PNGs. Cached via manifest.json.""" manifest_path = os.path.join(cache_dir, "manifest.json") # Include slice count in key to detect stale caches (e.g. different split) cache_key = f"{len(slices)}_{len(patient_ids)}" if os.path.exists(manifest_path): with open(manifest_path) as f: m = json.load(f) if m.get("_key") == cache_key and len(m["paths"]) == len(slices): return m["paths"], np.array(m["patient_ids"]), np.array(m["z_indices"]) else: print(f" Cache stale (key mismatch), re-exporting ...", flush=True) print(f" Exporting {len(slices)} PNGs to {cache_dir} ...", flush=True) paths, pids, zs = [], [], [] for i, (sl, pid, fname) in enumerate(zip(slices, patient_ids, filenames)): out_path = os.path.join(cache_dir, f"{fname}.png") Image.fromarray(sl).save(out_path) paths.append(out_path) pids.append(pid) zs.append(int(fname.split("_slice")[-1])) if (i + 1) % 500 == 0: print(f" {i + 1}/{len(slices)} ...", flush=True) with open(manifest_path, "w") as f: json.dump({"_key": f"{len(slices)}_{len(patient_ids)}", "paths": paths, "patient_ids": [str(p) for p in pids], "z_indices": [int(z) for z in zs]}, f) return paths, np.array(pids), np.array(zs) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): ap = argparse.ArgumentParser() ap.add_argument("--backbone", default="resnet18", choices=list(BACKBONES.keys())) ap.add_argument("--n-pairs", type=int, default=N_PAIRS) ap.add_argument("--epochs", type=int, default=EPOCHS) ap.add_argument("--lr", type=float, default=LR) ap.add_argument("--batch-size", type=int, default=BATCH_SIZE) ap.add_argument("--cache-dir", default=DEFAULT_CACHE_DIR) ap.add_argument("--flip-vertical", action="store_true", help="Flip slices vertically to match IQ-OTH orientation.") ap.add_argument("--rotate", type=int, default=0, choices=[0, 90, 180, 270], help="Rotate slices by N degrees (e.g. 90 if spine is on left).") ap.add_argument("--test-patients", type=int, default=TEST_PATIENTS, help="Number of patients held out for final testing.") ap.add_argument("--device", default=None) args = ap.parse_args() feat_dim, input_size = BACKBONES[args.backbone] device = torch.device(args.device or ( "cuda" if torch.cuda.is_available() else "cpu")) print(f"Device: {device}, backbone: {args.backbone}, " f"input_size: {input_size}") volume_dir = os.path.join(os.path.dirname(ROOT), "Task06_Lung", "imagesTr") # ---- Split patients into train / held-out test ---- # First, get the list of patients without loading all data from classes.nifti_dataset import NiftiSliceDataset as _DS import glob as _glob _nii = _glob.glob(os.path.join(volume_dir, "*.nii")) _gz = [p for p in _glob.glob(os.path.join(volume_dir, "*.nii.gz")) if os.path.basename(p)[:-7] not in {os.path.basename(q)[:-4] for q in _nii}] all_volumes = sorted(_nii + _gz) all_patient_ids = [os.path.basename(p).replace(".nii.gz", "").replace(".nii", "") for p in all_volumes] n_total = len(all_patient_ids) rng = np.random.default_rng(SEED) test_pids = set(rng.choice(all_patient_ids, size=min(args.test_patients, n_total - 1), replace=False)) train_pids_set = set(all_patient_ids) - test_pids print(f"\nPatients: {n_total} total → {len(train_pids_set)} train, " f"{len(test_pids)} held-out test\n") # ---- Extract TRAIN slices (sparse, for training speed) ---- print("=" * 50) print("Loading TRAIN slices (stride=3)") print("=" * 50) ds_train = _DS(volume_dir, random_slices=False, seed=SEED, flip_vertical=args.flip_vertical, rotate_deg=args.rotate) ds_train.load_all_slices(stride=3) # Filter to train patients only train_slice_mask = np.isin(ds_train.patient_labels, list(train_pids_set)) train_slices = [sl for i, sl in enumerate(ds_train.slices_rgb) if train_slice_mask[i]] train_pids_arr = ds_train.patient_labels[train_slice_mask] train_fnames = ds_train.filenames[train_slice_mask] train_zs = np.array([int(f.split("_slice")[-1]) for f in train_fnames]) # Export train PNGs train_cache = os.path.join(args.cache_dir, "train") os.makedirs(train_cache, exist_ok=True) train_paths, _, _ = _export_subset( train_slices, train_pids_arr, train_fnames, train_cache) print(f" Train: {len(train_paths)} slices, " f"{len(np.unique(train_pids_arr))} patients\n") # ---- Extract TEST slices (dense, for thorough eval) ---- print("=" * 50) print("Loading TEST slices (stride=1, full central 60%)") print("=" * 50) ds_test = _DS(volume_dir, random_slices=False, seed=SEED, flip_vertical=args.flip_vertical) ds_test.load_all_slices(stride=1) test_slice_mask = np.isin(ds_test.patient_labels, list(test_pids)) test_slices = [sl for i, sl in enumerate(ds_test.slices_rgb) if test_slice_mask[i]] test_pids_arr = ds_test.patient_labels[test_slice_mask] test_fnames = ds_test.filenames[test_slice_mask] test_zs = np.array([int(f.split("_slice")[-1]) for f in test_fnames]) # Export test PNGs test_cache = os.path.join(args.cache_dir, "test") os.makedirs(test_cache, exist_ok=True) test_paths, _, _ = _export_subset( test_slices, test_pids_arr, test_fnames, test_cache) print(f" Test: {len(test_paths)} slices, " f"{len(np.unique(test_pids_arr))} patients\n") # ---- Transforms ---- train_tf = transforms.Compose([ transforms.Resize((input_size, input_size)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) val_tf = transforms.Compose([ transforms.Resize((input_size, input_size)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # ---- Training pairs (80% of train patients for training, 20% for val monitoring) ---- train_unique = np.unique(train_pids_arr) n_val_patients = max(2, int(len(train_unique) * 0.2)) val_pids_set = set(rng.choice(train_unique, size=n_val_patients, replace=False)) tr_pids_set = set(train_unique) - val_pids_set tr_mask = np.isin(train_pids_arr, list(tr_pids_set)) val_mask = np.isin(train_pids_arr, list(val_pids_set)) tr_paths_list = [train_paths[i] for i in np.where(tr_mask)[0]] val_paths_list = [train_paths[i] for i in np.where(val_mask)[0]] print(f" Training pairs from: {len(tr_paths_list)} slices, " f"{len(tr_pids_set)} patients") print(f" Val pairs from: {len(val_paths_list)} slices, " f"{len(val_pids_set)} patients") print(f"\nGenerating {args.n_pairs:,} training pairs ...", flush=True) t0 = time.time() tr_ds = PairDataset(tr_paths_list, train_pids_arr[tr_mask], train_zs[tr_mask], n_pairs=args.n_pairs, transform=train_tf, seed=SEED) print(f" ... done ({time.time() - t0:.0f}s)") val_n = args.n_pairs // 4 print(f"Generating {val_n:,} validation pairs ...", flush=True) val_ds = PairDataset(val_paths_list, train_pids_arr[val_mask], train_zs[val_mask], n_pairs=val_n, transform=val_tf, seed=SEED + 1) tr_loader = DataLoader(tr_ds, batch_size=args.batch_size, shuffle=True, num_workers=4) val_loader = DataLoader(val_ds, batch_size=args.batch_size, num_workers=4) # ---- Model ---- model = SiameseCNN(args.backbone).to(device) print(f"\nModel: {sum(p.numel() for p in model.parameters()):,} parameters") for p in model.backbone.parameters(): p.requires_grad = False criterion = nn.BCEWithLogitsLoss() optimizer = optim.Adam(model.parameters(), lr=args.lr) # ---- Train ---- best_val_acc = 0 unfreeze_epoch = max(1, args.epochs // 2) print(f"\n{'Epoch':>6s} {'tr_loss':>8s} {'tr_acc':>8s} " f"{'val_loss':>8s} {'val_acc':>8s} {'val_auc':>8s}") print("-" * 52) for epoch in range(1, args.epochs + 1): if epoch == unfreeze_epoch + 1: for p in model.backbone.parameters(): p.requires_grad = True for g in optimizer.param_groups: g["lr"] = args.lr * 0.1 print(" (unfreezing backbone, LR 0.1x)", flush=True) tr_loss, tr_acc = train_epoch(model, tr_loader, optimizer, criterion, device) val_loss, val_acc, val_auc = eval_epoch(model, val_loader, criterion, device) star = "*" if val_acc > best_val_acc else "" if val_acc > best_val_acc: best_val_acc = val_acc os.makedirs(MODELS_DIR, exist_ok=True) torch.save(model.state_dict(), os.path.join(MODELS_DIR, f"siamese_{args.backbone}.pt")) print(f"{epoch:>6d} {tr_loss:>8.4f} {tr_acc:>8.4f} " f"{val_loss:>8.4f} {val_acc:>8.4f} {val_auc:>8.4f} {star}", flush=True) print(f"\nBest val accuracy: {best_val_acc:.4f}") # ---- Final evaluation on held-out TEST set ---- print(f"\n{'='*50}") print(f"Held-out TEST evaluation ({len(test_pids)} unseen patients, " f"{len(test_paths)} slices)") print("=" * 50) model_path = os.path.join(MODELS_DIR, f"siamese_{args.backbone}.pt") model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True)) model.eval() from classes.siamese import SiamesePatientMatcher as _SPM matcher = _SPM(model_path, backbone=args.backbone, device=str(device), input_size=input_size) test_pids_list = list(test_pids) test_fnames_clean = np.array([f.split("/")[-1].replace(".png", "") for f in test_paths]) # Connected components (no known K) manifest_cc = matcher.identify_patients( test_paths, filenames=test_fnames_clean, threshold=0.9, top_k=20, k=None) # Spectral clustering (known K) manifest_sc = matcher.identify_patients( test_paths, filenames=test_fnames_clean, threshold=0.9, top_k=20, k=len(test_pids)) def score_manifest(manifest, y_true_map): """Score a manifest against ground-truth patient IDs.""" # y_true_map: filename → true_patient_id all_labels = [] all_preds = [] for pred_pid, fnames in manifest.items(): for f in fnames: all_preds.append(pred_pid) all_labels.append(y_true_map.get(f, f"unknown_{f}")) all_labels = np.array(all_labels) all_preds = np.array(all_preds) purities = [] for c in np.unique(all_preds): mask = all_preds == c _, cts = np.unique(all_labels[mask], return_counts=True) purities.append(cts.max() / mask.sum()) captures = [] for p in np.unique(all_labels): mask = all_labels == p _, cts = np.unique(all_preds[mask], return_counts=True) captures.append(cts.max() / mask.sum()) return np.median(purities), np.median(captures), len(np.unique(all_preds)) # Build ground-truth map: filename → patient_id true_map = {f: p for f, p in zip(test_fnames_clean, test_pids_arr)} p_cc, c_cc, n_cc = score_manifest(manifest_cc, true_map) p_sc, c_sc, n_sc = score_manifest(manifest_sc, true_map) print(f"\n {'Method':<30s} {'Groups':>7s} {'Purity':>8s} {'Capture':>8s}") print(f" {'-'*53}") print(f" {'Connected components':<30s} {n_cc:>7d} " f"{p_cc:>7.1%} {c_cc:>7.1%}") print(f" {'Spectral (known K=' + str(len(test_pids)) + ')':<30s} {n_sc:>7d} " f"{p_sc:>7.1%} {c_sc:>7.1%}") # Compare to thumbnail baseline results_path = os.path.join(RESULTS_DIR, "task06_clustering_validation.json") if os.path.exists(results_path): with open(results_path) as f: prev = json.load(f) tp = prev["methods"]["thumbnail_64x64"]["cluster_purity"]["median"] tc = prev["methods"]["thumbnail_64x64"]["dominant_capture"]["median"] print(f" {'Thumbnail 64x64 (baseline)':<30s} {'—':>7s} " f"{tp:>7.1%} {tc:>7.1%}") print(f"\nModel saved → {model_path}") print("DONE") if __name__ == "__main__": main()