2026001
This commit is contained in:
rpotter6298
+86
-17
@@ -96,8 +96,11 @@ class FeatureExtractor:
|
||||
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)
|
||||
# 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":
|
||||
@@ -106,10 +109,16 @@ class FeatureExtractor:
|
||||
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
|
||||
# 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()
|
||||
@@ -162,25 +171,17 @@ class FeatureExtractor:
|
||||
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:
|
||||
for path in image_paths[start:end]:
|
||||
try:
|
||||
img = Image.open(path).convert("RGB")
|
||||
tensor = self.transform(img)
|
||||
batch_tensors.append(tensor)
|
||||
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))
|
||||
|
||||
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())
|
||||
features.append(self._forward(batch_tensors))
|
||||
|
||||
if (start // batch_size) % 20 == 0:
|
||||
print(f" Processed {end}/{n_images} images...")
|
||||
@@ -192,6 +193,74 @@ class FeatureExtractor:
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user