model.py

"""
model.py — CNN architecture reconstruction for chest X-ray classification.

Reconstructs the custom CNN matching the best_medical_model.pth state dict.
Architecture: Stem → 12 Feature Blocks → Classifier (3 classes).
"""

import torch
import torch.nn as nn
from pathlib import Path

CLASS_NAMES = ["Normal", "Pneumonia", "Tuberculosis"]
MODEL_PATH = Path(__file__).parent / "best_medical_model.pth"


class FeatureBlock(nn.Module):
    """Single feature block with two conv-bn-relu + two depthwise-conv-bn-relu paths."""

    def __init__(self, channels: int = 64):
        super().__init__()
        self.block = nn.Sequential(
            # 0: Conv2d(64, 64, 3x3)
            nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=True),
            # 1: BatchNorm2d(64)
            nn.BatchNorm2d(channels),
            # 2: ReLU
            nn.ReLU(inplace=False),
            # 3: Depthwise Conv2d(64, 1, 3x3, groups=64)
            nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels, bias=True),
            # 4: BatchNorm2d(64)
            nn.BatchNorm2d(channels),
            # 5: ReLU
            nn.ReLU(inplace=False),
            # 6: Conv2d(64, 64, 3x3)
            nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=True),
            # 7: BatchNorm2d(64)
            nn.BatchNorm2d(channels),
            # 8: ReLU
            nn.ReLU(inplace=False),
            # 9: ReLU (second activation before depthwise)
            nn.ReLU(inplace=False),
            # 10: Depthwise Conv2d(64, 1, 3x3, groups=64)
            nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels, bias=True),
            # 11: BatchNorm2d(64)
            nn.BatchNorm2d(channels),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.block(x)


class MedicalCNN(nn.Module):
    """Custom CNN for chest X-ray classification (Normal / Pneumonia / Tuberculosis)."""

    def __init__(self, num_classes: int = 3, channels: int = 64):
        super().__init__()
        self.stem = nn.Conv2d(3, channels, kernel_size=7, padding=3, bias=True)
        self.features = nn.Sequential(*[FeatureBlock(channels) for _ in range(12)])
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(channels, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        x = self.features(x)
        x = self.classifier(x)
        return x


def load_model(model_path: str = None, device: str = "cpu") -> MedicalCNN:
    """Load the trained medical CNN model from checkpoint."""
    path = Path(model_path) if model_path else MODEL_PATH
    model = MedicalCNN()
    state_dict = torch.load(str(path), map_location=device, weights_only=False)
    model.load_state_dict(state_dict)
    model.eval()
    model.to(device)
    return model