ryan/patient_leakage_detection
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

2026001

Browse Source
This commit is contained in:
rpotter6298
2026-07-01 17:35:58 +02:00
parent 9bfcc0243b
commit 35cbd9ac3c
84 changed files with 8500 additions and 423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
classes/siamese.py
+599
View File
@@ -0,0 +1,599 @@
"""
SiamesePatientMatcher — learned patient identification via connected components.
Trains a siamese CNN on Task06_Lung (known patient IDs), then applies the
trained model to IQ-OTH/NCCD to build a patient manifest without K-means.
Architecture:
Two CT slices → shared backbone → [f_A, f_B, |f_A-f_B|] → MLP head → same/different
For inference, we build a graph: edge between slices i,j if the siamese
confidence exceeds a threshold, then find connected components.
Each component = one estimated patient.
"""
import os
import numpy as np
from collections import defaultdict
import torch
import torch.nn as nn
from torchvision import transforms, models
# ---------------------------------------------------------------------------
# Backbone registry
# ---------------------------------------------------------------------------
BACKBONES = {
"resnet18": (models.resnet18, models.ResNet18_Weights.IMAGENET1K_V1, 512),
"resnet34": (models.resnet34, models.ResNet34_Weights.IMAGENET1K_V1, 512),
"efficientnet_b0": (models.efficientnet_b0, models.EfficientNet_B0_Weights.IMAGENET1K_V1, 1280),
}
# ---------------------------------------------------------------------------
# Model definition
# ---------------------------------------------------------------------------
class SiameseCNN(nn.Module):
"""Shared CNN backbone → concatenate → MLP head → binary classification."""
def __init__(self, backbone_name="resnet18", hidden_dims=None):
super().__init__()
if hidden_dims is None:
hidden_dims = [512, 128]
fn, weights, feat_dim = BACKBONES[backbone_name]
cnn = fn(weights=weights)
if backbone_name.startswith("resnet"):
self.backbone = nn.Sequential(
cnn.conv1, cnn.bn1, cnn.relu, cnn.maxpool,
cnn.layer1, cnn.layer2, cnn.layer3, cnn.layer4,
nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten())
elif backbone_name.startswith("efficientnet"):
self.backbone = nn.Sequential(
cnn.features, nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten())
else:
raise ValueError(f"Unknown backbone: {backbone_name}")
# Head: input = backbone_dim * 3
head_input = feat_dim * 3
layers = []
prev = head_input
for h in hidden_dims:
layers.append(nn.Linear(prev, h))
layers.append(nn.BatchNorm1d(h))
layers.append(nn.ReLU())
layers.append(nn.Dropout(0.3))
prev = h
layers.append(nn.Linear(prev, 1))
self.head = nn.Sequential(*layers)
def forward(self, img_a, img_b):
fa = self.backbone(img_a)
fb = self.backbone(img_b)
combined = torch.cat([fa, fb, torch.abs(fa - fb)], dim=-1)
return self.head(combined).squeeze(-1)
def embed(self, images, device):
"""Extract backbone feature vectors for a batch of images."""
return self.backbone(images)
# ---------------------------------------------------------------------------
# Patient matcher
# ---------------------------------------------------------------------------
class SiamesePatientMatcher:
"""Apply a trained siamese model to identify patient groups in new data.
Parameters
----------
model_path : str
Path to saved model weights (.pt file).
backbone : str
Backbone name matching the saved model.
device : str or None
Torch device. Auto-detected if None.
input_size : int
Input resolution for the backbone (224 for ResNet, 240 for EfficientNet).
"""
def __init__(self, model_path, backbone="resnet18", device=None,
input_size=224):
self.device = torch.device(
device or ("cuda" if torch.cuda.is_available() else "cpu"))
self.input_size = input_size
self.model = SiameseCNN(backbone).to(self.device)
self.model.load_state_dict(
torch.load(model_path, map_location=self.device, weights_only=True))
self.model.eval()
self.transform = 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]),
])
def identify_patients(self, images, filenames=None, threshold=0.9,
top_k=20, batch_size=64, k=None,
cluster_method="edge_rank",
min_size=None, max_size=None, keep_k=False):
"""Build a patient manifest.
If ``k`` is None (default): connected-components clustering with a
hard similarity threshold. No prior knowledge of patient count needed.
If ``k`` is provided: spectral clustering into exactly ``k`` groups
using the siamese similarity graph. Weak/spurious edges get cut to
respect the known patient count. Use this when you know K a priori
(e.g. IQ-OTH has 15 Benign, 40 Malignant, 55 Normal patients).
Parameters
----------
images : list
PIL Images, numpy arrays, or paths to image files.
filenames : list of str or None
Image filenames for the output manifest.
threshold : float
Min siamese probability for an edge (connected-components mode only).
top_k : int
Top-K similar candidates to verify per slice.
batch_size : int
Batch size for embedding extraction.
k : int or None
If provided, partition into exactly k groups via spectral clustering.
Returns
-------
dict mapping patient_id → list of filenames.
"""
if filenames is None:
filenames = [str(i) for i in range(len(images))]
# ---- Load and embed all images ----
print(f"Embedding {len(images)} images ...", flush=True)
all_embeddings = self.embed_images(images, batch_size)
print(f" Embeddings: {all_embeddings.shape}", flush=True)
# ---- Build similarity graph ----
sim_matrix, verified_edges = self._build_similarity_graph(
all_embeddings, threshold, top_k)
# ---- Cluster ----
if k is not None:
if cluster_method in ("complete", "average"):
# Agglomerative clustering on the dense siamese P(same) matrix.
# Complete/average linkage resist the single-linkage "chaining"
# that edge-ranking/connected-components suffer when the model
# is over-confident (e.g. out-of-distribution on IQ-OTH), where
# a few cross-patient edges merge many patients into one blob.
manifest = self._agglomerative_cluster(
all_embeddings, filenames, k, linkage=cluster_method)
else:
# Default: edge-ranking (preserves natural clusters), falling
# back to spectral if not enough edges to reach k.
manifest = self._edge_rank_cluster(
verified_edges, len(filenames), filenames, k)
if manifest is None:
print(" Edge-ranking couldn't reach k, falling back to "
"spectral", flush=True)
manifest = self._spectral_cluster(
sim_matrix, all_embeddings, filenames, k)
else:
manifest = self._connected_components(
verified_edges, len(filenames), filenames)
if min_size or max_size:
if keep_k and k is not None:
manifest = self._rebalance_keep_k(
manifest, filenames, all_embeddings, k, min_size, max_size)
else:
manifest = self._rebalance(
manifest, filenames, all_embeddings, min_size, max_size)
return manifest
# ------------------------------------------------------------------
# Internals
# ------------------------------------------------------------------
def _build_similarity_graph(self, embeddings, threshold, top_k):
"""Build a sparse similarity graph verified by the siamese head.
Returns
-------
sim_matrix : (n, n) ndarray cosine similarity (all pairs)
edges : list of (i, j, prob) verified edges above threshold
"""
n = len(embeddings)
# Cosine similarity
emb_norm = embeddings / (np.linalg.norm(embeddings, axis=1,
keepdims=True) + 1e-8)
sim_matrix = emb_norm @ emb_norm.T
top_k_idx = np.argsort(-sim_matrix, axis=1)[:, 1:top_k + 1]
# Verify top-K candidates with siamese head
print(f"Verifying top-{top_k} candidates per slice "
f"(threshold={threshold}) ...", flush=True)
edges = []
emb_t = torch.from_numpy(embeddings.astype(np.float32)).to(self.device)
for i in range(n):
for j in top_k_idx[i]:
if j <= i:
continue
with torch.no_grad():
fa = emb_t[i:i + 1]
fb = emb_t[j:j + 1]
combined = torch.cat([fa, fb, torch.abs(fa - fb)], dim=-1)
prob = torch.sigmoid(self.model.head(combined)).item()
if prob > threshold:
edges.append((i, j, prob))
if (i + 1) % 500 == 0:
print(f" {i + 1}/{n} slices, {len(edges)} edges",
flush=True)
print(f" Total verified edges: {len(edges)}", flush=True)
return sim_matrix, edges
def _connected_components(self, edges, n, filenames):
"""Cluster via connected components on verified edges."""
from scipy.sparse.csgraph import connected_components
from scipy.sparse import csr_matrix
if edges:
row, col = zip(*[(e[0], e[1]) for e in edges])
# Symmetric: each edge goes both ways, double the data
all_row = list(row) + list(col)
all_col = list(col) + list(row)
data = np.ones(len(all_row))
adj = csr_matrix((data, (all_row, all_col)), shape=(n, n))
n_components, labels = connected_components(adj, directed=False)
else:
n_components, labels = n, np.arange(n)
return self._labels_to_manifest(labels, n_components, filenames,
"siamese")
def embed_images(self, images, batch_size=64):
"""Backbone feature vectors for paths / PIL images / arrays."""
from PIL import Image as PILImage
embs = []
for start in range(0, len(images), batch_size):
batch = []
for item in images[start:start + batch_size]:
if isinstance(item, str):
img = PILImage.open(item).convert("RGB")
elif isinstance(item, PILImage.Image):
img = item.convert("RGB")
else:
arr = np.asarray(item)
if arr.dtype != np.uint8:
arr = (255.0 * (arr - arr.min()) /
(arr.max() - arr.min() + 1e-8)).astype(np.uint8)
img = PILImage.fromarray(arr).convert("RGB")
batch.append(self.transform(img))
batch_t = torch.stack(batch).to(self.device)
with torch.no_grad():
embs.append(self.model.backbone(batch_t).cpu().numpy())
if (start // batch_size) % 20 == 0:
print(f" {start + len(batch)}/{len(images)}", flush=True)
return np.concatenate(embs, axis=0)
def _dense_prob_matrix(self, embeddings):
"""Full symmetric P(same-patient) matrix from the siamese head.
Scores every pair with the head (cat[fa, fb, |fa-fb|] -> sigmoid),
one query row at a time to avoid materialising all n^2 vectors, then
symmetrises since the head is not exactly order-invariant.
"""
n = len(embeddings)
emb = torch.from_numpy(embeddings.astype(np.float32)).to(self.device)
P = np.zeros((n, n), dtype=np.float32)
with torch.no_grad():
for i in range(n):
fa = emb[i:i + 1].expand(n, -1)
combined = torch.cat([fa, emb, torch.abs(fa - emb)], dim=-1)
P[i] = torch.sigmoid(
self.model.head(combined).squeeze(-1)).cpu().numpy()
return 0.5 * (P + P.T)
def _agglomerative_cluster(self, embeddings, filenames, k, linkage="complete"):
"""Cluster into k groups via agglomerative clustering on 1 - P(same).
Unlike single-linkage (edge-ranking / connected components), complete
and average linkage will not merge two groups on the strength of a
single confident cross-patient edge, so they resist chaining.
"""
from sklearn.cluster import AgglomerativeClustering
print(f" Agglomerative clustering ({linkage} linkage) into k={k} ...",
flush=True)
P = self._dense_prob_matrix(embeddings)
dist = 1.0 - P
np.fill_diagonal(dist, 0.0)
dist[dist < 0] = 0.0
try:
model = AgglomerativeClustering(
n_clusters=k, metric="precomputed", linkage=linkage)
except TypeError: # scikit-learn < 1.2
model = AgglomerativeClustering(
n_clusters=k, affinity="precomputed", linkage=linkage)
labels = model.fit_predict(dist)
return self._labels_to_manifest(labels, k, filenames, "siamese")
def _rebalance(self, manifest, filenames, embeddings,
min_size=None, max_size=None):
"""Enforce group-size bounds using the siamese distances.
Oversized groups (> max_size) are split at their natural gaps via
complete-linkage on the members' 1 - P(same) submatrix — appropriate
when a group is really several distinct patients chained together.
Undersized groups (< min_size, e.g. orphaned singletons that are a
leakage hazard) are absorbed into their nearest group. Splitting first,
then absorbing, keeps sizes within bounds where possible.
"""
from sklearn.cluster import AgglomerativeClustering
f2i = {f: i for i, f in enumerate(filenames)}
dist = 1.0 - self._dense_prob_matrix(embeddings)
np.fill_diagonal(dist, 0.0)
dist[dist < 0] = 0.0
groups = {pid: [f2i[f] for f in fs if f in f2i]
for pid, fs in manifest.items()}
# 1) Split oversized groups. Complete-linkage into ceil(size/max_size)
# sub-groups can still leave one child over the cap (uneven splits),
# so repeat until every group is <= max_size.
if max_size:
changed = True
while changed:
changed = False
for pid in list(groups):
idx = groups[pid]
if len(idx) <= max_size:
continue
n_sub = int(np.ceil(len(idx) / max_size))
sub = dist[np.ix_(idx, idx)]
try:
model = AgglomerativeClustering(
n_clusters=n_sub, metric="precomputed", linkage="complete")
except TypeError:
model = AgglomerativeClustering(
n_clusters=n_sub, affinity="precomputed", linkage="complete")
lab = model.fit_predict(sub)
del groups[pid]
for s in range(n_sub):
members = [idx[j] for j in range(len(idx)) if lab[j] == s]
if members:
groups[f"{pid}_s{s}"] = members
changed = True
# 2) Absorb undersized groups into their nearest remaining group.
if min_size:
changed = True
while changed:
changed = False
for pid in sorted((p for p in groups if len(groups[p]) < min_size),
key=lambda p: len(groups[p])):
if pid not in groups or len(groups) == 1:
continue
idx = groups[pid]
best, best_d = None, np.inf
for opid, oidx in groups.items():
if opid == pid:
continue
d = dist[np.ix_(idx, oidx)].min()
if d < best_d:
best_d, best = d, opid
if best is not None:
groups[best] = groups[best] + idx
del groups[pid]
changed = True
return {pid: [filenames[i] for i in idx] for pid, idx in groups.items()}
def _rebalance_keep_k(self, manifest, filenames, embeddings, k,
min_size=None, max_size=None, max_iter=1000):
"""Rebalance group sizes while keeping exactly k groups.
Each pass splits the largest group into two balanced halves — bisecting
k-means on the siamese embeddings, which cuts at the natural density gap
and favours an even split rather than shaving off one or two points —
and merges the smallest group into its nearest neighbour, so the group
count is preserved. Repeats until every group is within
[min_size, max_size]. It is a no-op when no group violates the bounds
(e.g. Task06, where edge-ranking already gives clean per-patient sizes).
"""
from sklearn.cluster import KMeans
f2i = {f: i for i, f in enumerate(filenames)}
dist = 1.0 - self._dense_prob_matrix(embeddings)
np.fill_diagonal(dist, 0.0)
dist[dist < 0] = 0.0
clusters = [[f2i[f] for f in fs if f in f2i] for fs in manifest.values()]
def bisect(idx):
lab = KMeans(n_clusters=2, n_init=10, random_state=42).fit_predict(
embeddings[idx])
a = [idx[i] for i in range(len(idx)) if lab[i] == 0]
b = [idx[i] for i in range(len(idx)) if lab[i] == 1]
if not a or not b: # degenerate: even halves
half = len(idx) // 2
a, b = idx[:half], idx[half:]
return a, b
def merge_smallest():
si = min(range(len(clusters)), key=lambda i: len(clusters[i]))
small = clusters.pop(si)
ji = min(range(len(clusters)),
key=lambda j: dist[np.ix_(small, clusters[j])].min())
clusters[ji].extend(small)
# Start from exactly k groups (edge-rank already yields k; be safe).
while len(clusters) > k:
merge_smallest()
while len(clusters) < k:
bi = max(range(len(clusters)), key=lambda i: len(clusters[i]))
clusters += list(bisect(clusters.pop(bi)))
hi = max_size if max_size else float("inf")
lo = min_size if min_size else 0
for _ in range(max_iter):
sizes = [len(c) for c in clusters]
if max(sizes) <= hi and min(sizes) >= lo:
break
bi = max(range(len(clusters)), key=lambda i: len(clusters[i]))
clusters += list(bisect(clusters.pop(bi))) # +1 group
merge_smallest() # -1 group -> keeps k
return {f"siamese_{i}": [filenames[j] for j in c]
for i, c in enumerate(clusters)}
def _edge_rank_cluster(self, edges, n, filenames, k):
"""Cluster into k groups by adding edges in descending confidence order.
Starts with n isolated nodes, adds edges from most to least confident
until exactly k connected components form. Natural clusters stay intact;
only the weakest-split cluster gets divided.
Returns None if we can't reach k components.
"""
from scipy.sparse.csgraph import connected_components
from scipy.sparse import csr_matrix
if len(edges) == 0:
return None
# Sort edges by probability descending
sorted_edges = sorted(edges, key=lambda e: e[2], reverse=True)
# Binary search for the threshold that gives exactly k components
# Start with all edges → see how many components
all_row = [e[0] for e in sorted_edges] + [e[1] for e in sorted_edges]
all_col = [e[1] for e in sorted_edges] + [e[0] for e in sorted_edges]
all_data = np.ones(len(all_row))
adj_full = csr_matrix((all_data, (all_row, all_col)), shape=(n, n))
n_min, _ = connected_components(adj_full, directed=False)
if n_min > k:
print(f" Edge ranking: even with all {len(edges)} edges, "
f"only {n_min} components (need {k})", flush=True)
return None
if n_min == k:
# Perfect — all edges give exactly k components
_, labels = connected_components(adj_full, directed=False)
return self._labels_to_manifest(labels, k, filenames, "siamese")
# Binary search: find edge index where components == k
lo, hi = 0, len(sorted_edges)
best_labels = None
while lo < hi:
mid = (lo + hi) // 2
# Build graph with first `mid` edges
sub_edges = sorted_edges[:mid]
row = [e[0] for e in sub_edges] + [e[1] for e in sub_edges]
col = [e[1] for e in sub_edges] + [e[0] for e in sub_edges]
data = np.ones(len(row))
adj = csr_matrix((data, (row, col)), shape=(n, n))
n_comp, labels = connected_components(adj, directed=False)
if n_comp > k:
lo = mid + 1 # need more edges
elif n_comp < k:
hi = mid # too many edges
else:
best_labels = labels
hi = mid # try to find the earliest edge that achieves k
if best_labels is None:
return None
# Report the confidence at the split point
split_conf = sorted_edges[lo - 1][2] if lo > 0 else 1.0
print(f" Edge ranking: k={k} reached at confidence={split_conf:.4f} "
f"(edge {lo}/{len(sorted_edges)})", flush=True)
return self._labels_to_manifest(best_labels, k, filenames, "siamese")
def _spectral_cluster(self, sim_matrix, embeddings, filenames, k):
"""Cluster into exactly k groups via spectral clustering on the
siamese similarity graph.
Builds a weighted adjacency matrix from cosine similarity, then uses
spectral clustering (normalized cut) to partition into k groups.
Weak/spurious edges get cut to respect the known group count.
"""
from sklearn.cluster import SpectralClustering
from scipy.sparse import csr_matrix
n = len(embeddings)
# Build sparse weighted adjacency from top similarities
# Use top_k=50 for denser graph (spectral clustering needs connectivity)
top_k_dense = min(50, n - 1)
top_idx = np.argsort(-sim_matrix, axis=1)[:, 1:top_k_dense + 1]
row, col, data = [], [], []
for i in range(n):
for j in top_idx[i]:
if j <= i:
continue
# Weight = cosine similarity (in [0,1] after ReLU)
w = max(0.0, float(sim_matrix[i, j]))
if w > 0:
row.append(i); col.append(j); data.append(w)
row.append(j); col.append(i); data.append(w)
adj = csr_matrix((data, (row, col)), shape=(n, n))
print(f" Adjacency: {len(data) // 2} edges (top-{top_k_dense})",
flush=True)
# Spectral clustering
print(f" Spectral clustering into k={k} groups ...", flush=True)
sc = SpectralClustering(
n_clusters=k, affinity="precomputed",
random_state=42, n_init=20,
assign_labels="kmeans") # kmeans discretization is more stable
labels = sc.fit_predict(adj.toarray())
# If spectral clustering fails (disconnected graph), fall back to
# adding a small epsilon to connect components
n_unique = len(np.unique(labels))
if n_unique < k:
print(f" Warning: only {n_unique}/{k} clusters found; "
f"graph may be disconnected. Adding background connectivity.",
flush=True)
# Add weak background edges
adj_dense = adj.toarray()
adj_dense += 0.001 * (1.0 - np.eye(n))
labels = SpectralClustering(
n_clusters=k, affinity="precomputed",
random_state=42, n_init=20,
assign_labels="kmeans").fit_predict(adj_dense)
return self._labels_to_manifest(labels, k, filenames, "spectral")
def _labels_to_manifest(self, labels, n_groups, filenames, prefix):
"""Convert flat cluster labels to a manifest dict."""
components = defaultdict(list)
for i, c in enumerate(labels):
components[int(c)].append(filenames[i])
manifest = {}
for comp_id, comp_files in sorted(components.items()):
manifest[f"{prefix}_{comp_id:03d}"] = sorted(comp_files)
print(f" {len(manifest)} groups from {len(filenames)} slices",
flush=True)
sizes = [len(v) for v in manifest.values()]
if sizes:
print(f" Group sizes: min={min(sizes)}, max={max(sizes)}, "
f"mean={np.mean(sizes):.1f}", flush=True)
return manifest