rpotter6298
204 lines
7.5 KiB
Python
204 lines
7.5 KiB
Python
"""
|
||
Patient identification via K-means clustering.
|
||
|
||
Two modes:
|
||
- Thumbnail-based: cluster on 64×64 grayscale images (Andreas's approach).
|
||
- Feature-based: cluster on PCA-reduced CNN features.
|
||
|
||
Since the IQ-OTH/NCCD dataset does not provide patient IDs,
|
||
we estimate which images belong to the same patient by clustering
|
||
within each class.
|
||
"""
|
||
|
||
import os
|
||
import numpy as np
|
||
from PIL import Image
|
||
from sklearn.decomposition import PCA
|
||
from sklearn.cluster import KMeans
|
||
|
||
|
||
class PatientIdentifier:
|
||
"""Estimate patient groups via K-means clustering within each class.
|
||
|
||
Parameters
|
||
----------
|
||
patient_estimates : dict
|
||
Mapping from class_name -> number_of_patients, e.g.
|
||
{"Bengin cases": 15, "Malignant cases": 40, "Normal cases": 55}.
|
||
cluster_image_size : tuple of int
|
||
Grayscale thumbnail size for clustering (default (64, 64)).
|
||
random_state : int
|
||
Seed for PCA and K-means reproducibility.
|
||
"""
|
||
|
||
def __init__(self, patient_estimates, cluster_image_size=(64, 64),
|
||
random_state=42):
|
||
self.patient_estimates = patient_estimates
|
||
self.cluster_image_size = cluster_image_size
|
||
self.random_state = random_state
|
||
self.groups_ = None
|
||
self.labels_ = None
|
||
self.assignments_ = None # dict: class_name -> list of cluster_ids per sample
|
||
|
||
# ------------------------------------------------------------------
|
||
# Public API
|
||
# ------------------------------------------------------------------
|
||
|
||
def identify_from_thumbnails(self, base_path, filenames, labels,
|
||
valid_extensions=None):
|
||
"""Cluster images within each class using grayscale thumbnails.
|
||
|
||
(Original approach from Andreas — model-agnostic.)
|
||
"""
|
||
self.labels_ = labels
|
||
groups = self._run_clustering(
|
||
labels, filenames,
|
||
data_loader_fn=lambda class_name, idx: self._load_thumbnails(
|
||
base_path, class_name, filenames[idx]),
|
||
valid_extensions=valid_extensions,
|
||
)
|
||
self.groups_ = groups
|
||
return groups
|
||
|
||
def identify_from_features(self, X, labels):
|
||
"""Cluster images within each class using pre-extracted CNN features.
|
||
|
||
Parameters
|
||
----------
|
||
X : np.ndarray (n_images, n_features)
|
||
Feature matrix from a CNN (e.g. VGG16 output).
|
||
labels : np.ndarray (n_images,)
|
||
Class label strings in the same order as X.
|
||
|
||
Returns
|
||
-------
|
||
groups : np.ndarray (n_images,) group label strings
|
||
"""
|
||
self.labels_ = labels
|
||
groups = self._run_clustering(
|
||
labels, None,
|
||
data_loader_fn=lambda class_name, idx_class: X[idx_class],
|
||
)
|
||
self.groups_ = groups
|
||
return groups
|
||
|
||
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
|
||
if class_prefixes is None:
|
||
class_prefixes = {c: self._slugify(c)
|
||
for c in self.patient_estimates}
|
||
assignment = {}
|
||
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 = []
|
||
cluster_ids = []
|
||
for g in groups[idx_class]:
|
||
cid = int(g.split("_cluster_")[-1])
|
||
if cid not in seen:
|
||
seen.append(cid)
|
||
cluster_ids.append(seen.index(cid))
|
||
|
||
for local_id in seen:
|
||
patient_id = f"{prefix}_{local_id}"
|
||
mask = np.array(cluster_ids) == local_id
|
||
assignment[patient_id] = filenames[idx_class][mask].tolist()
|
||
|
||
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.")
|
||
return np.unique(self.groups_)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Internals
|
||
# ------------------------------------------------------------------
|
||
|
||
def _run_clustering(self, labels, filenames, data_loader_fn,
|
||
valid_extensions=None):
|
||
"""Shared K-means pipeline used by both thumbnail and feature modes."""
|
||
if valid_extensions is None:
|
||
valid_extensions = (".png", ".jpg", ".jpeg",
|
||
".tif", ".tiff", ".bmp")
|
||
|
||
n_samples = len(labels)
|
||
groups = np.array([None] * n_samples, dtype=object)
|
||
|
||
for class_name, n_clusters in self.patient_estimates.items():
|
||
print(f"\n Clustering class: {class_name} "
|
||
f"(k={n_clusters})")
|
||
|
||
idx_class = np.where(labels == class_name)[0]
|
||
|
||
X_cluster = data_loader_fn(class_name, idx_class)
|
||
if isinstance(X_cluster, list):
|
||
X_cluster = np.array(X_cluster, dtype=np.float32)
|
||
|
||
# Dimensionality reduction
|
||
n_pca = min(50, X_cluster.shape[0] - 1, X_cluster.shape[1])
|
||
pca = PCA(n_components=n_pca, random_state=self.random_state)
|
||
X_pca = pca.fit_transform(X_cluster)
|
||
|
||
# K-means
|
||
kmeans = KMeans(
|
||
n_clusters=n_clusters,
|
||
random_state=self.random_state,
|
||
n_init=20
|
||
)
|
||
clusters = kmeans.fit_predict(X_pca)
|
||
|
||
for i, cluster_id in zip(idx_class, clusters):
|
||
groups[i] = f"{class_name}_cluster_{cluster_id}"
|
||
|
||
cluster_sizes = np.bincount(clusters)
|
||
print(f" Cluster sizes: min={cluster_sizes.min()}, "
|
||
f"max={cluster_sizes.max()}, "
|
||
f"mean={cluster_sizes.mean():.1f}")
|
||
|
||
n_unique = len(np.unique(groups))
|
||
print(f"\n Estimated {n_unique} patient groups total")
|
||
return groups
|
||
|
||
def _load_thumbnails(self, base_path, class_name, files):
|
||
"""Load grayscale 64×64 thumbnails for a class."""
|
||
class_path = os.path.join(base_path, class_name)
|
||
X = []
|
||
for file in files:
|
||
img_path = os.path.join(class_path, file)
|
||
try:
|
||
img = Image.open(img_path).convert("L")
|
||
img = img.resize(self.cluster_image_size)
|
||
img = np.array(img, dtype=np.float32) / 255.0
|
||
X.append(img.flatten())
|
||
except Exception as e:
|
||
print(f" Error loading {file}: {e}")
|
||
X.append(np.zeros(
|
||
self.cluster_image_size[0] *
|
||
self.cluster_image_size[1],
|
||
dtype=np.float32
|
||
))
|
||
return np.array(X, dtype=np.float32)
|