rpotter6298
600 lines
25 KiB
Python
600 lines
25 KiB
Python
"""
|
|
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
|