""" 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"]