""" 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