rpotter6298
500 lines
19 KiB
Python
500 lines
19 KiB
Python
#!/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()
|