src/classifier_torch.py

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