rpotter6298
282 lines
10 KiB
Python
282 lines
10 KiB
Python
"""
|
||
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"]
|