2026001
This commit is contained in:
rpotter6298
@@ -1,3 +1,5 @@
|
||||
from .features import FeatureExtractor
|
||||
from .patient_identifier import PatientIdentifier
|
||||
from .classifier import PatientLeakageClassifier
|
||||
from .nifti_dataset import NiftiSliceDataset
|
||||
from .siamese import SiamesePatientMatcher, SiameseCNN
|
||||
|
||||
+106
-50
@@ -61,7 +61,7 @@ class PatientLeakageClassifier:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def run(self, model_name, seed, split_type="patient",
|
||||
nfeatures_list=None, gamma_logspace=None):
|
||||
nfeatures_list=None, gamma_logspace=None, return_predictions=False):
|
||||
"""Run full pipeline: train/test split → RF importance →
|
||||
GridSearchCV over nfeatures×gamma → retrain → test.
|
||||
|
||||
@@ -72,6 +72,9 @@ class PatientLeakageClassifier:
|
||||
split_type : str "image" or "patient".
|
||||
nfeatures_list : list Feature counts to search over.
|
||||
gamma_logspace : tuple Args for np.logspace (start, stop, n).
|
||||
return_predictions : bool
|
||||
If True, also include the test labels/predictions ("y_true",
|
||||
"y_pred") so callers can build a confusion matrix without refitting.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -84,55 +87,12 @@ class PatientLeakageClassifier:
|
||||
gamma_logspace = DEFAULT_GAMMA_RANGE
|
||||
|
||||
X, Y, groups = self._load_data(model_name)
|
||||
(X_tr, X_te, y_tr, y_te, g_tr, cv,
|
||||
n_tr_patients, n_te_patients) = self._split(X, Y, groups, seed, split_type)
|
||||
|
||||
# --- Train/test split ---
|
||||
if split_type == "image":
|
||||
X_tr, X_te, y_tr, y_te = train_test_split(
|
||||
X, Y, test_size=0.2, random_state=seed, stratify=Y)
|
||||
cv = 5
|
||||
g_tr = None
|
||||
n_tr_patients = None
|
||||
n_te_patients = None
|
||||
else:
|
||||
unique_g = np.unique(groups)
|
||||
g_labels = np.array([Y[groups == g][0] for g in unique_g])
|
||||
tr_g, te_g = train_test_split(
|
||||
unique_g, test_size=0.2, random_state=seed, stratify=g_labels)
|
||||
tr_mask = np.isin(groups, tr_g)
|
||||
te_mask = np.isin(groups, te_g)
|
||||
X_tr, X_te = X[tr_mask], X[te_mask]
|
||||
y_tr, y_te = Y[tr_mask], Y[te_mask]
|
||||
g_tr = groups[tr_mask]
|
||||
n_tr_patients = len(np.unique(g_tr))
|
||||
n_te_patients = len(np.unique(groups[te_mask]))
|
||||
cv = StratifiedGroupKFold(n_splits=5, shuffle=True,
|
||||
random_state=seed)
|
||||
|
||||
# --- RF feature importance ---
|
||||
fs = RandomForestClassifier(n_estimators=500, random_state=15,
|
||||
class_weight="balanced", n_jobs=self.n_jobs)
|
||||
fs.fit(X_tr, y_tr)
|
||||
sorted_idx = np.argsort(fs.feature_importances_)[::-1]
|
||||
|
||||
# --- Grid search over nfeatures × gamma ---
|
||||
pipe = Pipeline([("scaler", StandardScaler()),
|
||||
("svm", SVC(kernel="rbf", C=10))])
|
||||
valid_nfeats = [n for n in nfeatures_list if n <= X_tr.shape[1]]
|
||||
best_cv, best_n, best_g = 0, None, None
|
||||
|
||||
for nfeat in valid_nfeats:
|
||||
sel = sorted_idx[:nfeat]
|
||||
pg = {"svm__gamma": (1.0 / nfeat) * np.logspace(*gamma_logspace)}
|
||||
grid = GridSearchCV(pipe, pg, cv=cv, scoring="accuracy",
|
||||
n_jobs=self.n_jobs, verbose=0)
|
||||
if g_tr is not None:
|
||||
grid.fit(X_tr[:, sel], y_tr, groups=g_tr)
|
||||
else:
|
||||
grid.fit(X_tr[:, sel], y_tr)
|
||||
if grid.best_score_ > best_cv:
|
||||
best_cv = grid.best_score_
|
||||
best_n = nfeat
|
||||
best_g = grid.best_params_["svm__gamma"]
|
||||
sorted_idx = self._rf_ranking(X_tr, y_tr)
|
||||
best_cv, best_n, best_g = self._grid_search(
|
||||
X_tr, y_tr, g_tr, cv, sorted_idx, nfeatures_list, gamma_logspace)
|
||||
|
||||
# --- Retrain + test ---
|
||||
final = Pipeline([("scaler", StandardScaler()),
|
||||
@@ -140,7 +100,7 @@ class PatientLeakageClassifier:
|
||||
final.fit(X_tr[:, sorted_idx[:best_n]], y_tr)
|
||||
y_pred = final.predict(X_te[:, sorted_idx[:best_n]])
|
||||
|
||||
return {
|
||||
result = {
|
||||
"model": model_name,
|
||||
"seed": seed,
|
||||
"split_type": split_type,
|
||||
@@ -153,11 +113,107 @@ class PatientLeakageClassifier:
|
||||
"n_train_patients": n_tr_patients,
|
||||
"n_test_patients": n_te_patients,
|
||||
}
|
||||
if return_predictions:
|
||||
result["y_true"] = [str(v) for v in y_te]
|
||||
result["y_pred"] = [str(v) for v in y_pred]
|
||||
return result
|
||||
|
||||
def cv_curve(self, model_name, seed, split_type="image",
|
||||
nfeatures_list=None, gamma_logspace=None):
|
||||
"""CV accuracy as a function of the number of RF-selected features.
|
||||
|
||||
Same split → RF ranking → per-nfeatures γ-grid as ``run``, but returns
|
||||
the whole curve (used by the Figure 3 / S1 / S4 CV-vs-features plots)
|
||||
instead of only the best point.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict: {"curves": {nfeat: cv_acc}, "best_nfeat", "best_cv", "best_gamma"}
|
||||
"""
|
||||
if nfeatures_list is None:
|
||||
nfeatures_list = DEFAULT_NFEATURES
|
||||
if gamma_logspace is None:
|
||||
gamma_logspace = DEFAULT_GAMMA_RANGE
|
||||
|
||||
X, Y, groups = self._load_data(model_name)
|
||||
X_tr, _, y_tr, _, g_tr, cv, _, _ = self._split(
|
||||
X, Y, groups, seed, split_type)
|
||||
sorted_idx = self._rf_ranking(X_tr, y_tr)
|
||||
|
||||
pipe = Pipeline([("scaler", StandardScaler()),
|
||||
("svm", SVC(kernel="rbf", C=10))])
|
||||
valid_nfeats = [n for n in nfeatures_list if n <= X_tr.shape[1]]
|
||||
curves, best_cv, best_n, best_g = {}, 0, None, None
|
||||
for nfeat in valid_nfeats:
|
||||
pg = {"svm__gamma": (1.0 / nfeat) * np.logspace(*gamma_logspace)}
|
||||
grid = GridSearchCV(pipe, pg, cv=cv, scoring="accuracy",
|
||||
n_jobs=self.n_jobs, verbose=0)
|
||||
if g_tr is not None:
|
||||
grid.fit(X_tr[:, sorted_idx[:nfeat]], y_tr, groups=g_tr)
|
||||
else:
|
||||
grid.fit(X_tr[:, sorted_idx[:nfeat]], y_tr)
|
||||
curves[int(nfeat)] = float(grid.best_score_)
|
||||
if grid.best_score_ > best_cv:
|
||||
best_cv = float(grid.best_score_)
|
||||
best_n = int(nfeat)
|
||||
best_g = float(grid.best_params_["svm__gamma"])
|
||||
return {"curves": curves, "best_nfeat": best_n,
|
||||
"best_cv": best_cv, "best_gamma": best_g}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internals
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _split(self, X, Y, groups, seed, split_type):
|
||||
"""Image-level or patient-level 80:20 train/test split.
|
||||
|
||||
Returns (X_tr, X_te, y_tr, y_te, g_tr, cv, n_tr_patients, n_te_patients);
|
||||
g_tr is None and cv is 5 for image-level splits.
|
||||
"""
|
||||
if split_type == "image":
|
||||
X_tr, X_te, y_tr, y_te = train_test_split(
|
||||
X, Y, test_size=0.2, random_state=seed, stratify=Y)
|
||||
return X_tr, X_te, y_tr, y_te, None, 5, None, None
|
||||
|
||||
unique_g = np.unique(groups)
|
||||
g_labels = np.array([Y[groups == g][0] for g in unique_g])
|
||||
tr_g, te_g = train_test_split(
|
||||
unique_g, test_size=0.2, random_state=seed, stratify=g_labels)
|
||||
tr_mask = np.isin(groups, tr_g)
|
||||
te_mask = np.isin(groups, te_g)
|
||||
cv = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=seed)
|
||||
return (X[tr_mask], X[te_mask], Y[tr_mask], Y[te_mask],
|
||||
groups[tr_mask], cv, len(np.unique(groups[tr_mask])),
|
||||
len(np.unique(groups[te_mask])))
|
||||
|
||||
def _rf_ranking(self, X_tr, y_tr):
|
||||
"""Feature indices ranked by RF Gini importance (descending)."""
|
||||
fs = RandomForestClassifier(n_estimators=500, random_state=15,
|
||||
class_weight="balanced", n_jobs=self.n_jobs)
|
||||
fs.fit(X_tr, y_tr)
|
||||
return np.argsort(fs.feature_importances_)[::-1]
|
||||
|
||||
def _grid_search(self, X_tr, y_tr, g_tr, cv, sorted_idx,
|
||||
nfeatures_list, gamma_logspace):
|
||||
"""Search nfeatures × gamma; return (best_cv, best_nfeat, best_gamma)."""
|
||||
pipe = Pipeline([("scaler", StandardScaler()),
|
||||
("svm", SVC(kernel="rbf", C=10))])
|
||||
valid_nfeats = [n for n in nfeatures_list if n <= X_tr.shape[1]]
|
||||
best_cv, best_n, best_g = 0, None, None
|
||||
for nfeat in valid_nfeats:
|
||||
pg = {"svm__gamma": (1.0 / nfeat) * np.logspace(*gamma_logspace)}
|
||||
grid = GridSearchCV(pipe, pg, cv=cv, scoring="accuracy",
|
||||
n_jobs=self.n_jobs, verbose=0)
|
||||
if g_tr is not None:
|
||||
grid.fit(X_tr[:, sorted_idx[:nfeat]], y_tr, groups=g_tr)
|
||||
else:
|
||||
grid.fit(X_tr[:, sorted_idx[:nfeat]], y_tr)
|
||||
if grid.best_score_ > best_cv:
|
||||
best_cv = grid.best_score_
|
||||
best_n = nfeat
|
||||
best_g = grid.best_params_["svm__gamma"]
|
||||
return best_cv, best_n, best_g
|
||||
|
||||
def _load_data(self, model_name):
|
||||
data = np.load(os.path.join(self.features_dir,
|
||||
f"{model_name}_features.npz"),
|
||||
|
||||
+86
-17
@@ -96,8 +96,11 @@ class FeatureExtractor:
|
||||
self.model = full_model.features
|
||||
|
||||
elif name == "EfficientNetB1":
|
||||
# Keep conv stack, drop avgpool + classifier
|
||||
# → output (B, 1280, 7, 7) (at 224 px; 8×8 at 240 px)
|
||||
# Keep conv stack, drop avgpool + classifier.
|
||||
# Input 240×240 → 8×8 spatial (240/32 = 7.5 → 8 with padding).
|
||||
# Output: (B, 1280, 8, 8) → 81,920-d.
|
||||
# NOTE: manuscript incorrectly reports 62,720 (7×7×1280);
|
||||
# they used the 224px spatial size for the calculation.
|
||||
self.model = full_model.features
|
||||
|
||||
elif name == "MobileNetV2":
|
||||
@@ -106,10 +109,16 @@ class FeatureExtractor:
|
||||
self.model = full_model.features
|
||||
|
||||
elif name == "ResNet50":
|
||||
# Drop fc, keep everything *including* the adaptive avg pool
|
||||
# → output (B, 2048, 1, 1) — matches TF include_top=False
|
||||
full_model.fc = nn.Identity()
|
||||
self.model = full_model
|
||||
# Drop BOTH avgpool and FC → (B, 2048, 7, 7) → 100,352-d.
|
||||
# This matches TF ResNet50(include_top=False) with no pooling and
|
||||
# the manuscript's reported dimension, and keeps ResNet50 consistent
|
||||
# with the other four models (all use flattened spatial conv maps).
|
||||
# Empirically the spatial features give ~0.99 image-level accuracy,
|
||||
# vs ~0.97 for the GAP-pooled 2,048-d vector, which discards the
|
||||
# 7×7 grid the RF+SVM pipeline relies on.
|
||||
self.model = nn.Sequential(
|
||||
*list(full_model.children())[:-2]
|
||||
)
|
||||
|
||||
self.model.to(self.device)
|
||||
self.model.eval()
|
||||
@@ -162,25 +171,17 @@ class FeatureExtractor:
|
||||
features = []
|
||||
for start in range(0, n_images, batch_size):
|
||||
end = min(start + batch_size, n_images)
|
||||
batch_paths = image_paths[start:end]
|
||||
batch_tensors = []
|
||||
|
||||
for path in batch_paths:
|
||||
for path in image_paths[start:end]:
|
||||
try:
|
||||
img = Image.open(path).convert("RGB")
|
||||
tensor = self.transform(img)
|
||||
batch_tensors.append(tensor)
|
||||
batch_tensors.append(self.transform(img))
|
||||
except Exception as e:
|
||||
print(f" Error loading {path}: {e}")
|
||||
batch_tensors.append(torch.zeros(3, self.input_size, self.input_size))
|
||||
|
||||
batch = torch.stack(batch_tensors).to(self.device)
|
||||
|
||||
with torch.no_grad():
|
||||
batch_features = self.model(batch)
|
||||
batch_features = batch_features.view(batch_features.size(0), -1)
|
||||
|
||||
features.append(batch_features.cpu().numpy())
|
||||
features.append(self._forward(batch_tensors))
|
||||
|
||||
if (start // batch_size) % 20 == 0:
|
||||
print(f" Processed {end}/{n_images} images...")
|
||||
@@ -192,6 +193,74 @@ class FeatureExtractor:
|
||||
print(f" Feature matrix shape: {X.shape}")
|
||||
return X, Y, filenames
|
||||
|
||||
def extract_from_images(self, images, labels=None, filenames=None,
|
||||
batch_size=32):
|
||||
"""Extract features from in-memory images (dataset-layout agnostic).
|
||||
|
||||
Unlike ``extract``, this does not assume a directory tree of JPEGs
|
||||
grouped into class folders, so it works for any source that can
|
||||
produce image arrays (e.g. axial slices decoded from NIfTI volumes).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
images : sequence
|
||||
Each item may be a PIL.Image, an (H, W) grayscale array, or an
|
||||
(H, W, 3) RGB array (uint8 or float). Grayscale is expanded to
|
||||
RGB and everything is resized/normalised via ``self.transform``.
|
||||
labels, filenames : sequence or None
|
||||
Optional per-image class labels / names, returned unchanged so the
|
||||
output matches ``extract``'s (X, Y, filenames) contract.
|
||||
|
||||
Returns
|
||||
-------
|
||||
X : np.ndarray (n_images, n_features)
|
||||
Y : np.ndarray (n_images,)
|
||||
filenames : np.ndarray (n_images,)
|
||||
"""
|
||||
n_images = len(images)
|
||||
print(f" Extracting features from {n_images} in-memory images "
|
||||
f"({self.model_name})")
|
||||
|
||||
features = []
|
||||
for start in range(0, n_images, batch_size):
|
||||
end = min(start + batch_size, n_images)
|
||||
batch_tensors = [self.transform(self._to_pil(im))
|
||||
for im in images[start:end]]
|
||||
features.append(self._forward(batch_tensors))
|
||||
|
||||
if (start // batch_size) % 20 == 0:
|
||||
print(f" Processed {end}/{n_images} images...")
|
||||
|
||||
X = np.concatenate(features, axis=0).astype(np.float32)
|
||||
Y = np.array(labels) if labels is not None else np.array([None] * n_images)
|
||||
fnames = (np.array(filenames) if filenames is not None
|
||||
else np.arange(n_images))
|
||||
print(f" Feature matrix shape: {X.shape}")
|
||||
return X, Y, fnames
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internals
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _forward(self, batch_tensors):
|
||||
"""Run one batch of preprocessed tensors through the model."""
|
||||
batch = torch.stack(batch_tensors).to(self.device)
|
||||
with torch.no_grad():
|
||||
f = self.model(batch)
|
||||
f = f.view(f.size(0), -1)
|
||||
return f.cpu().numpy()
|
||||
|
||||
@staticmethod
|
||||
def _to_pil(im):
|
||||
"""Coerce a PIL image or numpy array (gray/RGB) to an RGB PIL image."""
|
||||
if isinstance(im, Image.Image):
|
||||
return im.convert("RGB")
|
||||
arr = np.asarray(im)
|
||||
if arr.dtype != np.uint8:
|
||||
lo, hi = float(arr.min()), float(arr.max())
|
||||
arr = (255.0 * (arr - lo) / (hi - lo + 1e-8)).astype(np.uint8)
|
||||
return Image.fromarray(arr).convert("RGB")
|
||||
|
||||
def save_features(self, X, Y, filenames, output_dir):
|
||||
"""Save extracted features to a compressed .npz file."""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
"""
|
||||
Load axial slices from a directory of NIfTI CT volumes.
|
||||
|
||||
Each ``.nii.gz`` file is treated as one patient (the filename stem is the
|
||||
ground-truth patient ID). A handful of evenly-spaced axial slices are sampled
|
||||
per volume, HU-windowed, and returned both as grayscale thumbnails (for the
|
||||
manuscript's clustering method) and as full-resolution slices (for CNN feature
|
||||
extraction).
|
||||
|
||||
This is used to validate the patient-clustering method on a dataset where the
|
||||
true patient identity IS known — e.g. the Medical Segmentation Decathlon
|
||||
Task06_Lung set — unlike IQ-OTH/NCCD, where patient IDs are unavailable.
|
||||
"""
|
||||
|
||||
import os
|
||||
import glob
|
||||
import numpy as np
|
||||
|
||||
|
||||
class NiftiSliceDataset:
|
||||
"""Sample HU-windowed axial slices from a folder of NIfTI CT volumes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
volume_dir : str
|
||||
Directory containing ``*.nii.gz`` volumes, one per patient.
|
||||
n_slices_per_patient : int
|
||||
Number of evenly-spaced axial slices to sample per volume.
|
||||
thumbnail_size : tuple of int
|
||||
Grayscale thumbnail size for clustering (matches the manuscript's 64x64).
|
||||
hu_window : tuple of float
|
||||
(low, high) Hounsfield-unit clip range applied before scaling to 8-bit.
|
||||
The default (-1000, 400) keeps the body outline, soft tissue, and lung
|
||||
parenchyma — the cues that identify a patient.
|
||||
central_fraction : float
|
||||
Fraction of the volume (centred on the middle slice) to sample from,
|
||||
avoiding empty/partial slices at the superior/inferior ends.
|
||||
slice_axis : int
|
||||
Axis along which axial slices are indexed (-1 = last, correct for the
|
||||
(H, W, Z) layout used by the Decathlon volumes).
|
||||
"""
|
||||
|
||||
def __init__(self, volume_dir, n_slices_per_patient=10,
|
||||
thumbnail_size=(64, 64), hu_window=(-1000.0, 400.0),
|
||||
central_fraction=0.6, slice_axis=-1, random_slices=True,
|
||||
seed=42, flip_vertical=False, rotate_deg=0):
|
||||
self.volume_dir = volume_dir
|
||||
self.n_slices_per_patient = n_slices_per_patient
|
||||
self.thumbnail_size = thumbnail_size
|
||||
self.hu_window = hu_window
|
||||
self.central_fraction = central_fraction
|
||||
self.slice_axis = slice_axis
|
||||
self.random_slices = random_slices
|
||||
self.seed = seed
|
||||
self.flip_vertical = flip_vertical
|
||||
self.rotate_deg = rotate_deg # 0, 90, 180, or 270
|
||||
|
||||
# Populated by load()
|
||||
self.thumbnails = None # (n_slices, H*W) float32 in [0, 1]
|
||||
self.slices_rgb = None # list of (H, W) uint8 arrays
|
||||
self.patient_labels = None # (n_slices,) ground-truth patient IDs
|
||||
self.filenames = None # (n_slices,) per-slice names
|
||||
|
||||
def load(self, verbose=True):
|
||||
"""Load and slice every volume in ``volume_dir``."""
|
||||
import nibabel as nib
|
||||
from PIL import Image
|
||||
|
||||
# Prefer decompressed .nii (fast memmap random-slice access) and fall
|
||||
# back to .nii.gz, skipping any .gz that has an uncompressed twin.
|
||||
nii = glob.glob(os.path.join(self.volume_dir, "*.nii"))
|
||||
stems = {os.path.basename(p)[:-4] for p in nii}
|
||||
gz = [p for p in glob.glob(os.path.join(self.volume_dir, "*.nii.gz"))
|
||||
if os.path.basename(p)[:-7] not in stems]
|
||||
paths = sorted(nii + gz)
|
||||
if not paths:
|
||||
raise FileNotFoundError(
|
||||
f"No .nii/.nii.gz volumes found in {self.volume_dir}")
|
||||
|
||||
thumbs, slices_rgb, patient_labels, filenames = [], [], [], []
|
||||
lo, hi = self.hu_window
|
||||
|
||||
for path in paths:
|
||||
patient_id = os.path.basename(path).replace(".nii.gz", "").replace(".nii", "")
|
||||
img = nib.load(path)
|
||||
n_z = img.shape[self.slice_axis]
|
||||
axis = self.slice_axis % img.ndim
|
||||
|
||||
for k in self._slice_indices(n_z):
|
||||
# Lazy slice: pull only this plane off disk instead of
|
||||
# materialising the whole float64 volume with get_fdata().
|
||||
slicer = [slice(None)] * img.ndim
|
||||
slicer[axis] = k
|
||||
sl = np.asarray(img.dataobj[tuple(slicer)], dtype=np.float32)
|
||||
# HU window -> 8-bit grayscale
|
||||
sl = np.clip(sl, lo, hi)
|
||||
sl = (255.0 * (sl - lo) / (hi - lo)).astype(np.uint8)
|
||||
if self.flip_vertical:
|
||||
sl = np.flipud(sl)
|
||||
if self.rotate_deg:
|
||||
sl = np.rot90(sl, k=self.rotate_deg // 90)
|
||||
|
||||
thumb = Image.fromarray(sl).convert("L").resize(
|
||||
self.thumbnail_size)
|
||||
thumbs.append(np.asarray(thumb, dtype=np.float32).flatten() / 255.0)
|
||||
slices_rgb.append(sl)
|
||||
patient_labels.append(patient_id)
|
||||
filenames.append(f"{patient_id}_slice{k:03d}")
|
||||
|
||||
if verbose:
|
||||
print(f" {patient_id}: {n_z} slices -> "
|
||||
f"sampled {self.n_slices_per_patient}")
|
||||
|
||||
self.thumbnails = np.array(thumbs, dtype=np.float32)
|
||||
self.slices_rgb = slices_rgb
|
||||
self.patient_labels = np.array(patient_labels)
|
||||
self.filenames = np.array(filenames)
|
||||
|
||||
if verbose:
|
||||
print(f"\n Loaded {len(self.filenames)} slices from "
|
||||
f"{len(paths)} patients")
|
||||
return self
|
||||
|
||||
def load_all_slices(self, stride=3, verbose=True):
|
||||
"""Load ALL slices from the central fraction (not just a sample).
|
||||
|
||||
Every ``stride``-th slice is taken to keep feature extraction manageable
|
||||
while still giving K-means enough data to cluster per patient.
|
||||
|
||||
Returns self (populates thumbnails, slices_rgb, patient_labels, filenames).
|
||||
"""
|
||||
import nibabel as nib
|
||||
from PIL import Image
|
||||
|
||||
nii = glob.glob(os.path.join(self.volume_dir, "*.nii"))
|
||||
stems = {os.path.basename(p)[:-4] for p in nii}
|
||||
gz = [p for p in glob.glob(os.path.join(self.volume_dir, "*.nii.gz"))
|
||||
if os.path.basename(p)[:-7] not in stems]
|
||||
paths = sorted(nii + gz)
|
||||
if not paths:
|
||||
raise FileNotFoundError(
|
||||
f"No .nii/.nii.gz volumes found in {self.volume_dir}")
|
||||
|
||||
thumbs, slices_rgb, patient_labels, filenames = [], [], [], []
|
||||
lo, hi = self.hu_window
|
||||
total_slices = 0
|
||||
|
||||
for path in paths:
|
||||
patient_id = os.path.basename(path).replace(".nii.gz", "").replace(".nii", "")
|
||||
img = nib.load(path)
|
||||
n_z = img.shape[self.slice_axis]
|
||||
axis = self.slice_axis % img.ndim
|
||||
|
||||
# All slices in central fraction, stepped
|
||||
half = self.central_fraction / 2.0
|
||||
start = int(n_z * (0.5 - half))
|
||||
end = int(n_z * (0.5 + half))
|
||||
indices = list(range(start, end, stride))
|
||||
|
||||
for k in indices:
|
||||
slicer = [slice(None)] * img.ndim
|
||||
slicer[axis] = k
|
||||
sl = np.asarray(img.dataobj[tuple(slicer)], dtype=np.float32)
|
||||
sl = np.clip(sl, lo, hi)
|
||||
sl = (255.0 * (sl - lo) / (hi - lo)).astype(np.uint8)
|
||||
if self.flip_vertical:
|
||||
sl = np.flipud(sl)
|
||||
if self.rotate_deg:
|
||||
sl = np.rot90(sl, k=self.rotate_deg // 90)
|
||||
|
||||
thumb = Image.fromarray(sl).convert("L").resize(
|
||||
self.thumbnail_size)
|
||||
thumbs.append(np.asarray(thumb, dtype=np.float32).flatten() / 255.0)
|
||||
slices_rgb.append(sl)
|
||||
patient_labels.append(patient_id)
|
||||
filenames.append(f"{patient_id}_slice{k:03d}")
|
||||
total_slices += 1
|
||||
|
||||
if verbose:
|
||||
print(f" {patient_id}: {n_z} slices -> {len(indices)} sampled")
|
||||
|
||||
self.thumbnails = np.array(thumbs, dtype=np.float32)
|
||||
self.slices_rgb = slices_rgb
|
||||
self.patient_labels = np.array(patient_labels)
|
||||
self.filenames = np.array(filenames)
|
||||
|
||||
if verbose:
|
||||
print(f"\n Loaded {total_slices} slices from "
|
||||
f"{len(paths)} patients (stride={stride})")
|
||||
return self
|
||||
|
||||
def export_pngs(self, cache_dir):
|
||||
"""Export loaded slices as PNG files for fast reloading.
|
||||
|
||||
Creates ``cache_dir/`` with one PNG per slice and a manifest.json
|
||||
mapping paths → patient IDs and z-indices.
|
||||
|
||||
Returns (paths, patient_ids, z_indices).
|
||||
"""
|
||||
import json
|
||||
from PIL import Image as PILImage
|
||||
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
manifest_path = os.path.join(cache_dir, "manifest.json")
|
||||
|
||||
if os.path.exists(manifest_path):
|
||||
print(f" Loading cached PNGs from {cache_dir}")
|
||||
with open(manifest_path) as f:
|
||||
manifest = json.load(f)
|
||||
return (manifest["paths"], np.array(manifest["patient_ids"]),
|
||||
np.array(manifest["z_indices"]))
|
||||
|
||||
if self.slices_rgb is None or self.patient_labels is None:
|
||||
raise RuntimeError("Call load() or load_all_slices() first.")
|
||||
|
||||
print(f" Exporting {len(self.slices_rgb)} PNGs to {cache_dir} ...")
|
||||
paths, pids, zs = [], [], []
|
||||
for i, (sl, pid, fname) in enumerate(
|
||||
zip(self.slices_rgb, self.patient_labels, self.filenames)):
|
||||
out_path = os.path.join(cache_dir, f"{fname}.png")
|
||||
PILImage.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(self.slices_rgb)} ...", flush=True)
|
||||
|
||||
manifest = {"paths": paths, "patient_ids": pids, "z_indices": zs}
|
||||
with open(manifest_path, "w") as f:
|
||||
json.dump(manifest, f)
|
||||
print(f" Exported {len(paths)} PNGs")
|
||||
return paths, np.array(pids), np.array(zs)
|
||||
|
||||
@property
|
||||
def n_patients(self):
|
||||
return len(np.unique(self.patient_labels))
|
||||
|
||||
def _slice_indices(self, n_z):
|
||||
"""Slice indices within the central fraction (even or random)."""
|
||||
half = self.central_fraction / 2.0
|
||||
start = int(n_z * (0.5 - half))
|
||||
end = int(n_z * (0.5 + half))
|
||||
if self.random_slices:
|
||||
rng = np.random.default_rng(self.seed)
|
||||
return np.sort(rng.integers(start, end, size=self.n_slices_per_patient))
|
||||
else:
|
||||
return np.linspace(start, end - 1, self.n_slices_per_patient).astype(int)
|
||||
@@ -53,8 +53,8 @@ class PatientIdentifier:
|
||||
self.labels_ = labels
|
||||
groups = self._run_clustering(
|
||||
labels, filenames,
|
||||
data_loader_fn=lambda class_name, files: self._load_thumbnails(
|
||||
base_path, class_name, files),
|
||||
data_loader_fn=lambda class_name, idx: self._load_thumbnails(
|
||||
base_path, class_name, filenames[idx]),
|
||||
valid_extensions=valid_extensions,
|
||||
)
|
||||
self.groups_ = groups
|
||||
@@ -82,21 +82,28 @@ class PatientIdentifier:
|
||||
self.groups_ = groups
|
||||
return groups
|
||||
|
||||
def build_assignment_dict(self, groups, filenames, labels):
|
||||
def build_assignment_dict(self, groups, filenames, labels,
|
||||
class_prefixes=None):
|
||||
"""Convert flat group array into a human-readable dict.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
class_prefixes : dict or None
|
||||
Optional mapping from class_name -> short prefix used in the
|
||||
patient IDs (e.g. {"Bengin cases": "benign"}). If None, a prefix
|
||||
is derived automatically from each class name, so the method works
|
||||
for any dataset rather than only IQ-OTH's three classes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict mapping patient_id (e.g. "benign_0") -> list of image filenames.
|
||||
"""
|
||||
self.labels_ = labels
|
||||
class_map = {
|
||||
"Bengin cases": "benign",
|
||||
"Malignant cases": "malig",
|
||||
"Normal cases": "normal",
|
||||
}
|
||||
if class_prefixes is None:
|
||||
class_prefixes = {c: self._slugify(c)
|
||||
for c in self.patient_estimates}
|
||||
assignment = {}
|
||||
for class_name, prefix in class_map.items():
|
||||
for class_name, prefix in class_prefixes.items():
|
||||
idx_class = np.where(labels == class_name)[0]
|
||||
# Get unique cluster IDs within this class (preserve order)
|
||||
seen = []
|
||||
@@ -115,6 +122,12 @@ class PatientIdentifier:
|
||||
self.assignments_ = assignment
|
||||
return assignment
|
||||
|
||||
@staticmethod
|
||||
def _slugify(class_name):
|
||||
"""Turn a class name into a short lowercase prefix for patient IDs."""
|
||||
token = str(class_name).strip().lower().split()[0]
|
||||
return "".join(ch for ch in token if ch.isalnum()) or "group"
|
||||
|
||||
def get_unique_groups(self):
|
||||
if self.groups_ is None:
|
||||
raise RuntimeError("Call identify_*() first.")
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user