""" 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. # 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": # Keep conv stack, drop adaptive pool + classifier # → output (B, 1280, 7, 7) self.model = full_model.features elif name == "ResNet50": # 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() # 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_tensors = [] for path in image_paths[start:end]: try: img = Image.open(path).convert("RGB") 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)) 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) filenames = np.array(fnames) 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) 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"]