Spaces:
Runtime error
Runtime error
File size: 8,309 Bytes
1cf3825 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | """PyTorch ports of the three classifiers originally written in TF/Keras.
Why a port: the upstream Neuro-Lens-AI repo's classifier weights ship as Git LFS
pointer files (134 bytes each), so a zip download from GitHub does not include
the actual binaries. The one .h5 that IS in upstream (real_eval_current/vit)
was trained with a different version of the model code (ResNet50 flattened at
the top level; our build_vit_classifier nests it as 'vit_hybrid_resnet_base')
and is topology-incompatible with src/models.py. TF 2.21 on this machine is
CPU-only (no native-Windows GPU), so CPU retraining is ~3 hrs per model.
PyTorch on the RTX 4060 cuts that to ~5 min per model. Architectures are
matched to src/models.py as closely as possible so behaviour matches what the
original paper describes.
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
class CNNClassifier(nn.Module):
"""Custom 3-block CNN baseline. Matches src/models.py:build_cnn_baseline.
3 Conv2D+MaxPool blocks (32 -> 64 -> 128 channels) -> Flatten -> Dense(128)
-> Dropout -> Dense(1, sigmoid). The TF version uses Rescaling(1/255) as the
first layer; we expect callers to pass images already in [0,1].
"""
def __init__(self, dropout: float = 0.3):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2),
nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2),
nn.Dropout(dropout),
)
# 224 / 2 / 2 / 2 = 28, so feature map is 128 x 28 x 28.
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128 * 28 * 28, 128),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(128, 1),
)
@property
def last_conv_module(self):
"""Grad-CAM target. Property avoids duplicating features[6] in state_dict."""
return self.features[6] # the Conv2d(64, 128, ...)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.classifier(self.features(x))
class TransferClassifier(nn.Module):
"""ResNet50 backbone (ImageNet pretrained) + classification head. Matches
src/models.py:build_transfer_model with default args (resnet50, fine_tune=False).
The backbone outputs (B, 2048, 7, 7); we GAP to (B, 2048), then
Dropout -> Dense(256, relu) -> Dropout -> Dense(1).
"""
def __init__(self, dropout: float = 0.3, freeze_backbone: bool = True):
super().__init__()
# weights=ResNet50_Weights.IMAGENET1K_V2 gives the better pretrained weights
self.backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
feature_dim = self.backbone.fc.in_features # 2048
# Replace the classification head with our own
self.backbone.fc = nn.Identity()
if freeze_backbone:
for p in self.backbone.parameters():
p.requires_grad = False
# Keep BN layers in eval mode regardless of model.train()
for m in self.backbone.modules():
if isinstance(m, nn.BatchNorm2d):
m.eval()
self.head = nn.Sequential(
nn.Dropout(dropout),
nn.Linear(feature_dim, 256),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(256, 1),
)
@property
def last_conv_module(self):
"""Grad-CAM target. Not registered as a child module via this property,
so the state_dict shape stays unchanged."""
return self.backbone.layer4[-1]
def forward(self, x: torch.Tensor) -> torch.Tensor:
features = self.backbone(x) # (B, 2048)
return self.head(features)
class ViTHybridClassifier(nn.Module):
"""Hybrid ResNet50 + 4 transformer-block ViT. Matches
src/models.py:build_vit_classifier with default args.
Frozen ResNet50 outputs (B, 2048, 7, 7) -> Conv1x1 to projection_dim (128)
-> reshape to (B, 49, 128) token sequence -> learnable position embedding
-> 4 transformer encoder blocks (4 heads, mlp_dim=256) -> LayerNorm ->
GlobalAveragePool over tokens -> Dropout -> Dense(128, relu) -> Dense(1).
"""
def __init__(self, projection_dim: int = 128, num_layers: int = 4, num_heads: int = 4,
mlp_dim: int = 256, dropout: float = 0.1):
super().__init__()
self.backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
# Remove the classification head and avgpool so we keep the feature map.
self.backbone.fc = nn.Identity()
self.backbone.avgpool = nn.Identity()
# Freeze
for p in self.backbone.parameters():
p.requires_grad = False
for m in self.backbone.modules():
if isinstance(m, nn.BatchNorm2d):
m.eval()
# 1x1 patch projection
self.patch_projection = nn.Conv2d(2048, projection_dim, kernel_size=1)
self.num_patches = 7 * 7 # 224 / 32 = 7
self.projection_dim = projection_dim
self.position_embedding = nn.Parameter(torch.zeros(1, self.num_patches, projection_dim))
nn.init.trunc_normal_(self.position_embedding, std=0.02)
encoder_layer = nn.TransformerEncoderLayer(
d_model=projection_dim,
nhead=num_heads,
dim_feedforward=mlp_dim,
dropout=dropout,
activation='gelu',
batch_first=True,
norm_first=False, # match TF behaviour: norm after, not before
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.final_norm = nn.LayerNorm(projection_dim)
self.head = nn.Sequential(
nn.Dropout(dropout),
nn.Linear(projection_dim, 128),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(128, 1),
)
@property
def last_conv_module(self):
"""Grad-CAM target on the ResNet50 last conv block. Property, not a
registered submodule, so state_dict layout matches the trained weights."""
return self.backbone.layer4[-1]
def _run_backbone(self, x: torch.Tensor) -> torch.Tensor:
"""ResNet50 forward up through layer4, keeping the (B, 2048, 7, 7) feature map."""
b = self.backbone
x = b.conv1(x); x = b.bn1(x); x = b.relu(x); x = b.maxpool(x)
x = b.layer1(x); x = b.layer2(x); x = b.layer3(x); x = b.layer4(x)
return x # (B, 2048, 7, 7)
def forward(self, x: torch.Tensor) -> torch.Tensor:
feat = self._run_backbone(x) # (B, 2048, 7, 7)
patches = self.patch_projection(feat) # (B, 128, 7, 7)
tokens = patches.flatten(2).transpose(1, 2) # (B, 49, 128)
tokens = tokens + self.position_embedding
tokens = self.transformer(tokens) # (B, 49, 128)
tokens = self.final_norm(tokens)
pooled = tokens.mean(dim=1) # (B, 128)
return self.head(pooled)
def get_classifier(model_name: str) -> nn.Module:
name = model_name.lower()
if name == 'cnn':
return CNNClassifier()
if name == 'transfer':
return TransferClassifier()
if name == 'vit':
return ViTHybridClassifier()
raise ValueError(f'Unknown classifier: {model_name!r}. Choose cnn / transfer / vit.')
@torch.no_grad()
def predict_probability(model: nn.Module, image_chw_01: torch.Tensor) -> float:
"""Run inference on a single image already in (C, H, W) layout, [0,1]
range. Returns the sigmoid probability of the 'tumor' class."""
model.eval()
logit = model(image_chw_01.unsqueeze(0))
return float(torch.sigmoid(logit).item())
__all__ = ['CNNClassifier', 'TransferClassifier', 'ViTHybridClassifier', 'get_classifier', 'predict_probability']
|