Add scripts for patient classification and visualization
- Created `classification.py` for comparing image-level and patient-level classification results using various CNN models. - Implemented `create_patient_groups.py` to extract features, generate PCA/t-SNE plots, and identify patient groups via K-means clustering. - Added `figure6.py` to generate boxplots for test accuracy across multiple seeds. - Developed `simple_patient_tsne.py` to perform t-SNE visualization of patient groups and save results in a manifest file. - Introduced `simple_patient_manifest.csv` to store patient IDs, classes, image counts, and associated images.
This commit is contained in:
rpotter6298
@@ -0,0 +1,212 @@
|
||||
"""
|
||||
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"]
|
||||
Reference in New Issue
Block a user