patient_leakage_detection/classes/patient_identifier.py

"""
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, files: self._load_thumbnails(
                base_path, class_name, files),
            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):
        """Convert flat group array into a human-readable dict.

        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",
        }
        assignment = {}
        for class_name, prefix in class_map.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

    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)