patient_leakage_detection/classes/features.py

"""
Feature extraction using pretrained CNNs (PyTorch).

Saves feature maps as .npz files with the model name in the filename,
e.g. "VGG16_features.npz" containing X, Y, and filenames arrays.

Supports all five models from the manuscript:
    VGG16, DenseNet121, EfficientNetB1, MobileNetV2, ResNet50
"""

import os
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import (
    vgg16, VGG16_Weights,
    densenet121, DenseNet121_Weights,
    efficientnet_b1, EfficientNet_B1_Weights,
    mobilenet_v2, MobileNet_V2_Weights,
    resnet50, ResNet50_Weights,
)


# ---------------------------------------------------------------------------
# Model registry
# ---------------------------------------------------------------------------

MODEL_CONFIGS = {
    "VGG16": {
        "fn": vgg16,
        "weights": VGG16_Weights.IMAGENET1K_V1,
        "input_size": 224,
    },
    "DenseNet121": {
        "fn": densenet121,
        "weights": DenseNet121_Weights.IMAGENET1K_V1,
        "input_size": 224,
    },
    "EfficientNetB1": {
        "fn": efficientnet_b1,
        "weights": EfficientNet_B1_Weights.IMAGENET1K_V1,
        "input_size": 240,
    },
    "MobileNetV2": {
        "fn": mobilenet_v2,
        "weights": MobileNet_V2_Weights.IMAGENET1K_V1,
        "input_size": 224,
    },
    "ResNet50": {
        "fn": resnet50,
        "weights": ResNet50_Weights.IMAGENET1K_V1,
        "input_size": 224,
    },
}


class FeatureExtractor:
    """Extract deep features from images using a pretrained CNN.

    Parameters
    ----------
    model_name : str
        One of: VGG16, DenseNet121, EfficientNetB1, MobileNetV2, ResNet50.
    device : str or None
        Torch device string. Auto-detected if None.
    """

    def __init__(self, model_name="VGG16", device=None):
        if model_name not in MODEL_CONFIGS:
            raise ValueError(
                f"Unsupported model '{model_name}'. "
                f"Choose from: {list(MODEL_CONFIGS.keys())}"
            )
        self.model_name = model_name
        self.cfg = MODEL_CONFIGS[model_name]
        self.input_size = self.cfg["input_size"]
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None
        self._load_model()

    def _load_model(self):
        """Load pretrained model and strip the classifier head."""
        full_model = self.cfg["fn"](weights=self.cfg["weights"])

        name = self.model_name

        if name == "VGG16":
            # Drop the classifier Sequential → output (B, 512, 7, 7)
            self.model = nn.Sequential(*list(full_model.children())[:-1])

        elif name == "DenseNet121":
            # Keep conv stack (dense blocks), drop classifier Linear
            # → output (B, 1024, 7, 7)
            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)
            self.model = full_model.features

        elif name == "MobileNetV2":
            # Keep conv stack, drop adaptive pool + classifier
            # → output (B, 1280, 7, 7)
            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

        self.model.to(self.device)
        self.model.eval()

        # Preprocessing
        self.transform = transforms.Compose([
            transforms.Resize((self.input_size, self.input_size)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            ),
        ])

    # ------------------------------------------------------------------
    # extract / save / load — unchanged public API
    # ------------------------------------------------------------------

    def extract(self, image_dir, valid_extensions=None, batch_size=32):
        """Extract features from all images in a directory tree.

        Expects subdirectories named by class (e.g. "Bengin cases/").

        Returns
        -------
        X : np.ndarray  (n_images, n_features)
        Y : np.ndarray  (n_images,)  class labels
        filenames : np.ndarray  (n_images,)  image filenames
        """
        if valid_extensions is None:
            valid_extensions = (".png", ".jpg", ".jpeg",
                                ".tif", ".tiff", ".bmp")

        image_paths, labels, fnames = [], [], []
        for class_name in sorted(os.listdir(image_dir)):
            class_path = os.path.join(image_dir, class_name)
            if not os.path.isdir(class_path):
                continue
            print(f"  Scanning class: {class_name}")
            for file in sorted(os.listdir(class_path)):
                if not file.lower().endswith(valid_extensions):
                    continue
                image_paths.append(os.path.join(class_path, file))
                labels.append(class_name)
                fnames.append(file)

        n_images = len(image_paths)
        print(f"  Found {n_images} images across all classes")

        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:
                try:
                    img = Image.open(path).convert("RGB")
                    tensor = self.transform(img)
                    batch_tensors.append(tensor)
                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())

            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)
        filenames = np.array(fnames)

        print(f"  Feature matrix shape: {X.shape}")
        return X, Y, filenames

    def save_features(self, X, Y, filenames, output_dir):
        """Save extracted features to a compressed .npz file."""
        os.makedirs(output_dir, exist_ok=True)
        filepath = os.path.join(output_dir, f"{self.model_name}_features.npz")
        np.savez_compressed(filepath, X=X, Y=Y, filenames=filenames)
        print(f"  Saved features to {filepath}")
        return filepath

    @staticmethod
    def load_features(filepath):
        """Load saved features from a .npz file.

        Returns
        -------
        X, Y, filenames : np.ndarray
        """
        data = np.load(filepath, allow_pickle=True)
        return data["X"], data["Y"], data["filenames"]