Auto-deploy from GitHub: 495db78a06be79166200269bb14d9e9b1e8906d6

Dockerfile

@@ -13,11 +13,12 @@ COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
 # Copy app files
-COPY app.py .
+COPY src/ src/
+COPY api/ api/
 COPY models/ models/
 
 # Expose port 7860 (HF Spaces default)
 EXPOSE 7860
 
 # Run the API
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "api.main:app", "--host", "0.0.0.0", "--port", "7860"]

api/Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    libgl1 \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for caching
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy app files
+COPY src/ src/
+COPY api/ api/
+COPY models/ models/
+
+# Expose port 7860 (HF Spaces default)
+EXPOSE 7860
+
+# Run the API
+CMD ["uvicorn", "api.main:app", "--host", "0.0.0.0", "--port", "7860"]

api/__init__.py

@@ -0,0 +1,9 @@
+"""FastAPI application for Pneumonia Detection."""
+
+from .main import app
+from .schemas import (
+    HealthResponse,
+    PredictionResponse,
+    GradCAMResponse,
+    ErrorResponse
+)

api/main.py

@@ -0,0 +1,245 @@
+"""
+FastAPI application for Pneumonia Detection API.
+
+Run with: uvicorn api.main:app --reload
+"""
+
+import io
+import time
+import base64
+from pathlib import Path
+
+import torch
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+from .schemas import (
+    HealthResponse,
+    PredictionResponse,
+    GradCAMResponse,
+    ErrorResponse
+)
+
+import sys
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.config import CHECKPOINT_PATH, CLASS_NAMES, CONFIDENCE_THRESHOLD
+from src.model import create_model, get_device
+from src.predict import load_model, predict_image
+from src.gradcam import generate_gradcam
+
+# =============================================================================
+# App Configuration
+# =============================================================================
+
+app = FastAPI(
+    title="Pneumonia Detection API",
+    description="Deep learning API for detecting pneumonia from chest X-ray images using EfficientNet-B0",
+    version="1.0.0",
+    docs_url="/docs",
+    redoc_url="/redoc"
+)
+
+# CORS middleware for frontend access
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Configure appropriately for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# =============================================================================
+# Model Loading (on startup)
+# =============================================================================
+
+model = None
+device = None
+
+
+@app.on_event("startup")
+async def load_model_on_startup():
+    """Load model when the API starts."""
+    global model, device
+
+    device = get_device()
+    print(f"Using device: {device}")
+
+    if not CHECKPOINT_PATH.exists():
+        print(f"Warning: Model checkpoint not found at {CHECKPOINT_PATH}")
+        return
+
+    model = create_model(pretrained=False, freeze_backbone=False, device=device)
+    model = load_model(model, CHECKPOINT_PATH, device)
+    print(f"Model loaded from {CHECKPOINT_PATH}")
+
+
+# =============================================================================
+# Helper Functions
+# =============================================================================
+
+ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png"}
+
+
+def validate_image(file: UploadFile) -> None:
+    """Validate uploaded image file."""
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid content type: {file.content_type}. Expected image/*"
+        )
+
+    ext = Path(file.filename).suffix.lower() if file.filename else ""
+    if ext not in ALLOWED_EXTENSIONS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid file extension: {ext}. Allowed: {ALLOWED_EXTENSIONS}"
+        )
+
+
+async def read_image(file: UploadFile) -> Image.Image:
+    """Read uploaded file as PIL Image."""
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        return image
+    except Exception as e:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Failed to read image: {str(e)}"
+        )
+
+
+# =============================================================================
+# API Endpoints
+# =============================================================================
+
+@app.get("/", include_in_schema=False)
+async def root():
+    """Redirect to docs."""
+    return {"message": "Pneumonia Detection API", "docs": "/docs"}
+
+
+@app.get("/health", response_model=HealthResponse, tags=["Health"])
+async def health_check():
+    """
+    Health check endpoint.
+
+    Returns the API status and model loading state.
+    """
+    return HealthResponse(
+        status="healthy" if model is not None else "model_not_loaded",
+        model_loaded=model is not None,
+        model_path=str(CHECKPOINT_PATH)
+    )
+
+
+@app.post(
+    "/predict",
+    response_model=PredictionResponse,
+    responses={400: {"model": ErrorResponse}, 503: {"model": ErrorResponse}},
+    tags=["Prediction"]
+)
+async def predict(file: UploadFile = File(..., description="Chest X-ray image (JPEG/PNG)")):
+    """
+    Predict pneumonia from chest X-ray image.
+
+    Upload a chest X-ray image and get the prediction (NORMAL or PNEUMONIA)
+    with confidence score.
+    """
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    validate_image(file)
+    image = await read_image(file)
+
+    # Run inference
+    start_time = time.time()
+    pred_class, confidence = predict_image(model, image, device)
+    processing_time = (time.time() - start_time) * 1000  # Convert to ms
+
+    # Calculate raw probability
+    probability = confidence if pred_class == "PNEUMONIA" else 1 - confidence
+
+    return PredictionResponse(
+        prediction=pred_class,
+        confidence=confidence,
+        probability=probability,
+        processing_time_ms=round(processing_time, 2)
+    )
+
+
+@app.post(
+    "/predict/gradcam",
+    response_model=GradCAMResponse,
+    responses={400: {"model": ErrorResponse}, 503: {"model": ErrorResponse}},
+    tags=["Prediction"]
+)
+async def predict_with_gradcam(file: UploadFile = File(..., description="Chest X-ray image (JPEG/PNG)")):
+    """
+    Predict with Grad-CAM visualization.
+
+    Returns prediction along with a Grad-CAM heatmap overlay showing
+    which regions of the image influenced the prediction.
+    """
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    validate_image(file)
+    image = await read_image(file)
+
+    # Run inference with Grad-CAM
+    start_time = time.time()
+    cam_image, pred_class, confidence, _ = generate_gradcam(model, image, device)
+    processing_time = (time.time() - start_time) * 1000
+
+    # Convert Grad-CAM image to base64
+    cam_pil = Image.fromarray(cam_image)
+    buffer = io.BytesIO()
+    cam_pil.save(buffer, format="PNG")
+    buffer.seek(0)
+    img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+    # Calculate raw probability
+    probability = confidence if pred_class == "PNEUMONIA" else 1 - confidence
+
+    return GradCAMResponse(
+        prediction=pred_class,
+        confidence=confidence,
+        probability=probability,
+        processing_time_ms=round(processing_time, 2),
+        gradcam_image=f"data:image/png;base64,{img_base64}"
+    )
+
+
+# =============================================================================
+# Error Handlers
+# =============================================================================
+
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request, exc):
+    """Handle HTTP exceptions."""
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={"error": exc.detail, "detail": None}
+    )
+
+
+@app.exception_handler(Exception)
+async def general_exception_handler(request, exc):
+    """Handle unexpected exceptions."""
+    return JSONResponse(
+        status_code=500,
+        content={"error": "Internal server error", "detail": str(exc)}
+    )
+
+
+# =============================================================================
+# Run with: uvicorn api.main:app --reload --host 0.0.0.0 --port 8000
+# =============================================================================
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

api/schemas.py

@@ -0,0 +1,72 @@
+"""
+Pydantic models for API request/response validation.
+"""
+
+from pydantic import BaseModel, Field
+from typing import Optional
+from enum import Enum
+
+
+class ClassLabel(str, Enum):
+    """Prediction class labels."""
+    NORMAL = "NORMAL"
+    PNEUMONIA = "PNEUMONIA"
+
+
+class HealthResponse(BaseModel):
+    """Health check response."""
+    status: str = Field(..., example="healthy")
+    model_loaded: bool = Field(..., example=True)
+    model_path: str = Field(..., example="models/best_model.pt")
+
+
+class PredictionResponse(BaseModel):
+    """Prediction response."""
+    prediction: ClassLabel = Field(..., description="Predicted class")
+    confidence: float = Field(..., ge=0, le=1, description="Confidence score")
+    probability: float = Field(..., ge=0, le=1, description="Raw probability for PNEUMONIA")
+    processing_time_ms: float = Field(..., description="Inference time in milliseconds")
+
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "prediction": "PNEUMONIA",
+                "confidence": 0.92,
+                "probability": 0.92,
+                "processing_time_ms": 45.2
+            }
+        }
+
+
+class GradCAMResponse(BaseModel):
+    """Prediction with Grad-CAM visualization."""
+    prediction: ClassLabel = Field(..., description="Predicted class")
+    confidence: float = Field(..., ge=0, le=1, description="Confidence score")
+    probability: float = Field(..., ge=0, le=1, description="Raw probability for PNEUMONIA")
+    processing_time_ms: float = Field(..., description="Inference time in milliseconds")
+    gradcam_image: str = Field(..., description="Base64 encoded Grad-CAM overlay image")
+
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "prediction": "PNEUMONIA",
+                "confidence": 0.92,
+                "probability": 0.92,
+                "processing_time_ms": 150.5,
+                "gradcam_image": "data:image/png;base64,..."
+            }
+        }
+
+
+class ErrorResponse(BaseModel):
+    """Error response."""
+    error: str = Field(..., description="Error message")
+    detail: Optional[str] = Field(None, description="Detailed error information")
+
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "error": "Invalid image format",
+                "detail": "Supported formats: JPEG, PNG"
+            }
+        }

requirements.txt

@@ -1,7 +1,45 @@
-torch>=2.0.0
-torchvision>=0.15.0
-fastapi>=0.100.0
-uvicorn>=0.23.0
-python-multipart>=0.0.6
+# Core Deep Learning
+torch>=2.1.0
+torchvision>=0.16.0
 pillow>=10.0.0
 numpy>=1.24.0
+
+# Data Analysis & Visualization
+pandas>=2.0.0
+matplotlib>=3.7.0
+seaborn>=0.12.0
+
+# Experiment Tracking
+wandb>=0.15.0
+
+# Model Interpretability
+grad-cam>=1.4.0
+
+# API
+fastapi>=0.104.0
+uvicorn>=0.24.0
+python-multipart>=0.0.6
+
+# Web UI
+streamlit>=1.28.0
+
+# Testing
+pytest>=7.4.0
+
+# Code Quality
+black>=23.0.0
+ruff>=0.1.0
+
+# Jupyter
+jupyterlab>=4.0.0
+ipywidgets>=8.0.0
+
+# Utilities
+python-dotenv>=1.0.0
+tqdm>=4.66.0
+scikit-learn>=1.3.0
+
+# ONNX Export
+onnx>=1.15.0
+onnxruntime>=1.16.0
+onnxscript>=0.1.0

src/__init__.py

@@ -0,0 +1,6 @@
+"""
+Pneumonia Detection from Chest X-Rays
+Medical Image Classification using Deep Learning
+"""
+
+__version__ = "0.1.0"

src/config.py

@@ -0,0 +1,112 @@
+"""
+Configuration constants for the Pneumonia Detection project.
+All hyperparameters and paths are defined here for easy modification.
+"""
+
+from pathlib import Path
+
+# =============================================================================
+# Project Paths
+# =============================================================================
+PROJECT_ROOT = Path(__file__).parent.parent
+DATA_DIR = PROJECT_ROOT / "data" / "raw"
+PROCESSED_DIR = PROJECT_ROOT / "data" / "processed"
+MODEL_DIR = PROJECT_ROOT / "models"
+OUTPUT_DIR = PROJECT_ROOT / "outputs"
+FIGURES_DIR = OUTPUT_DIR / "figures"
+LOGS_DIR = OUTPUT_DIR / "logs"
+
+# =============================================================================
+# Data Configuration
+# =============================================================================
+IMAGE_SIZE = 224  # EfficientNet-B0 input size
+BATCH_SIZE = 32
+NUM_WORKERS = 4  # DataLoader workers
+
+# ImageNet normalization (required for pretrained models)
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+# Class labels
+CLASS_NAMES = ["NORMAL", "PNEUMONIA"]
+NUM_CLASSES = 1  # Binary classification with sigmoid
+
+# =============================================================================
+# Model Configuration
+# =============================================================================
+MODEL_NAME = "efficientnet_b0"
+DROPOUT_RATE = 0.3
+PRETRAINED = True
+
+# =============================================================================
+# Training Configuration - Stage 1 (Frozen Backbone)
+# =============================================================================
+STAGE1_EPOCHS = 5
+STAGE1_LR = 1e-4
+STAGE1_FREEZE_BACKBONE = True
+
+# =============================================================================
+# Training Configuration - Stage 2 (Fine-tuning)
+# =============================================================================
+STAGE2_EPOCHS = 15
+STAGE2_LR = 1e-5
+STAGE2_FREEZE_BACKBONE = False
+
+# =============================================================================
+# Optimizer Configuration
+# =============================================================================
+WEIGHT_DECAY = 1e-4
+BETAS = (0.9, 0.999)
+
+# =============================================================================
+# Scheduler Configuration
+# =============================================================================
+SCHEDULER_PATIENCE = 3
+SCHEDULER_FACTOR = 0.5
+SCHEDULER_MIN_LR = 1e-7
+
+# =============================================================================
+# Early Stopping Configuration
+# =============================================================================
+EARLY_STOP_PATIENCE = 7
+EARLY_STOP_MIN_DELTA = 0.001
+
+# =============================================================================
+# Model Checkpointing
+# =============================================================================
+CHECKPOINT_PATH = MODEL_DIR / "best_model.pt"
+SAVE_BEST_ONLY = True
+MONITOR_METRIC = "val_loss"
+
+# =============================================================================
+# Weights & Biases Configuration
+# =============================================================================
+WANDB_PROJECT = "pneumonia-detection"
+WANDB_ENTITY = None  # Set to your W&B username if needed
+
+# =============================================================================
+# Inference Configuration
+# =============================================================================
+CONFIDENCE_THRESHOLD = 0.5  # For binary classification
+GRADCAM_TARGET_LAYER = "features"  # EfficientNet feature extractor
+
+# =============================================================================
+# Random Seed (for reproducibility)
+# =============================================================================
+SEED = 42
+
+
+def create_directories():
+    """Create all necessary directories if they don't exist."""
+    for directory in [DATA_DIR, PROCESSED_DIR, MODEL_DIR, FIGURES_DIR, LOGS_DIR]:
+        directory.mkdir(parents=True, exist_ok=True)
+
+
+if __name__ == "__main__":
+    # Print configuration for verification
+    print(f"Project Root: {PROJECT_ROOT}")
+    print(f"Data Directory: {DATA_DIR}")
+    print(f"Model Directory: {MODEL_DIR}")
+    print(f"Image Size: {IMAGE_SIZE}")
+    print(f"Batch Size: {BATCH_SIZE}")
+    print(f"Model: {MODEL_NAME}")

src/dataset.py

@@ -0,0 +1,201 @@
+"""
+PyTorch Dataset and DataLoader utilities for Chest X-Ray classification.
+"""
+
+from pathlib import Path
+from typing import Tuple, Optional, List
+import random
+
+import torch
+from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
+from torchvision import transforms
+from PIL import Image
+from sklearn.model_selection import train_test_split
+
+from .config import (
+    DATA_DIR, IMAGE_SIZE, BATCH_SIZE, NUM_WORKERS,
+    IMAGENET_MEAN, IMAGENET_STD, CLASS_NAMES, SEED
+)
+
+
+class ChestXRayDataset(Dataset):
+    """Dataset for Chest X-Ray images."""
+
+    def __init__(
+        self,
+        image_paths: List[Path],
+        labels: List[int],
+        transform: Optional[transforms.Compose] = None
+    ):
+        self.image_paths = image_paths
+        self.labels = labels
+        self.transform = transform
+
+    def __len__(self) -> int:
+        return len(self.image_paths)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
+        img_path = self.image_paths[idx]
+        label = self.labels[idx]
+
+        # Load image and convert to RGB
+        image = Image.open(img_path).convert('RGB')
+
+        if self.transform:
+            image = self.transform(image)
+
+        return image, label
+
+
+def get_transforms(is_training: bool = True) -> transforms.Compose:
+    """Get image transforms for training or validation/test."""
+    if is_training:
+        return transforms.Compose([
+            transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.RandomRotation(10),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+        ])
+    else:
+        return transforms.Compose([
+            transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
+        ])
+
+
+def load_image_paths_and_labels(
+    data_dir: Path,
+    split: str
+) -> Tuple[List[Path], List[int]]:
+    """Load image paths and labels from a data split directory."""
+    image_paths = []
+    labels = []
+
+    for class_idx, class_name in enumerate(CLASS_NAMES):
+        class_dir = data_dir / split / class_name
+        if class_dir.exists():
+            for img_path in class_dir.glob('*.jpeg'):
+                image_paths.append(img_path)
+                labels.append(class_idx)
+
+    return image_paths, labels
+
+
+def create_train_val_split(
+    data_dir: Path = DATA_DIR,
+    val_ratio: float = 0.15,
+    seed: int = SEED
+) -> Tuple[List[Path], List[int], List[Path], List[int]]:
+    """Create stratified train/val split from training data."""
+    # Load all training images
+    train_paths, train_labels = load_image_paths_and_labels(data_dir, 'train')
+
+    # Stratified split
+    train_paths, val_paths, train_labels, val_labels = train_test_split(
+        train_paths, train_labels,
+        test_size=val_ratio,
+        stratify=train_labels,
+        random_state=seed
+    )
+
+    return train_paths, train_labels, val_paths, val_labels
+
+
+def get_class_weights(labels: List[int]) -> torch.Tensor:
+    """Calculate class weights for imbalanced dataset."""
+    class_counts = torch.bincount(torch.tensor(labels))
+    total = len(labels)
+    weights = total / (len(class_counts) * class_counts.float())
+    return weights
+
+
+def get_sampler(labels: List[int]) -> WeightedRandomSampler:
+    """Create weighted sampler for balanced batches."""
+    class_weights = get_class_weights(labels)
+    sample_weights = [class_weights[label] for label in labels]
+    sampler = WeightedRandomSampler(
+        weights=sample_weights,
+        num_samples=len(labels),
+        replacement=True
+    )
+    return sampler
+
+
+def get_dataloaders(
+    data_dir: Path = DATA_DIR,
+    batch_size: int = BATCH_SIZE,
+    num_workers: int = NUM_WORKERS,
+    val_ratio: float = 0.15,
+    use_weighted_sampling: bool = True
+) -> Tuple[DataLoader, DataLoader, DataLoader]:
+    """Create train, validation, and test DataLoaders."""
+
+    # Create train/val split
+    train_paths, train_labels, val_paths, val_labels = create_train_val_split(
+        data_dir, val_ratio
+    )
+
+    # Load test data
+    test_paths, test_labels = load_image_paths_and_labels(data_dir, 'test')
+
+    # Create datasets
+    train_dataset = ChestXRayDataset(
+        train_paths, train_labels, transform=get_transforms(is_training=True)
+    )
+    val_dataset = ChestXRayDataset(
+        val_paths, val_labels, transform=get_transforms(is_training=False)
+    )
+    test_dataset = ChestXRayDataset(
+        test_paths, test_labels, transform=get_transforms(is_training=False)
+    )
+
+    # Create sampler for training if using weighted sampling
+    train_sampler = get_sampler(train_labels) if use_weighted_sampling else None
+
+    # Only use pin_memory for CUDA (not supported on MPS)
+    pin_memory = torch.cuda.is_available()
+
+    # Create dataloaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        sampler=train_sampler,
+        shuffle=(train_sampler is None),
+        num_workers=num_workers,
+        pin_memory=pin_memory
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=pin_memory
+    )
+
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=pin_memory
+    )
+
+    # Print dataset info
+    print(f"Train: {len(train_dataset)} images")
+    print(f"Val:   {len(val_dataset)} images")
+    print(f"Test:  {len(test_dataset)} images")
+
+    return train_loader, val_loader, test_loader
+
+
+def get_pos_weight(labels: List[int]) -> torch.Tensor:
+    """Calculate pos_weight for BCEWithLogitsLoss to handle class imbalance."""
+    labels_tensor = torch.tensor(labels)
+    neg_count = (labels_tensor == 0).sum().float()  # NORMAL
+    pos_count = (labels_tensor == 1).sum().float()  # PNEUMONIA
+    pos_weight = neg_count / pos_count
+    return pos_weight

src/evaluate.py

@@ -0,0 +1,107 @@
+"""
+Evaluation functions for Pneumonia classification.
+"""
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import numpy as np
+from typing import Dict, Tuple
+from sklearn.metrics import (
+    accuracy_score, precision_score, recall_score, f1_score,
+    roc_auc_score, confusion_matrix, classification_report
+)
+
+from .config import CLASS_NAMES
+
+
+def predict_proba(
+    model: nn.Module,
+    loader: DataLoader,
+    device: torch.device
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Get predictions, probabilities, and true labels."""
+    model.eval()
+    all_probs, all_preds, all_labels = [], [], []
+
+    with torch.no_grad():
+        for images, labels in loader:
+            images = images.to(device)
+            outputs = model(images)
+            probs = torch.sigmoid(outputs).cpu().numpy()
+            preds = (probs > 0.5).astype(int)
+
+            all_probs.extend(probs.flatten())
+            all_preds.extend(preds.flatten())
+            all_labels.extend(labels.numpy())
+
+    return np.array(all_probs), np.array(all_preds), np.array(all_labels)
+
+
+def compute_metrics(y_true: np.ndarray, y_pred: np.ndarray, y_proba: np.ndarray) -> Dict:
+    """Compute all evaluation metrics."""
+    return {
+        'accuracy': accuracy_score(y_true, y_pred),
+        'precision': precision_score(y_true, y_pred),
+        'recall': recall_score(y_true, y_pred),
+        'f1': f1_score(y_true, y_pred),
+        'roc_auc': roc_auc_score(y_true, y_proba),
+        'confusion_matrix': confusion_matrix(y_true, y_pred)
+    }
+
+
+def evaluate_model(
+    model: nn.Module,
+    loader: DataLoader,
+    device: torch.device
+) -> Dict:
+    """Full evaluation on a dataset."""
+    probs, preds, labels = predict_proba(model, loader, device)
+    metrics = compute_metrics(labels, preds, probs)
+
+    print("=" * 50)
+    print("EVALUATION RESULTS")
+    print("=" * 50)
+    print(f"Accuracy:  {metrics['accuracy']:.4f}")
+    print(f"Precision: {metrics['precision']:.4f}")
+    print(f"Recall:    {metrics['recall']:.4f}")
+    print(f"F1 Score:  {metrics['f1']:.4f}")
+    print(f"ROC-AUC:   {metrics['roc_auc']:.4f}")
+    print("\nConfusion Matrix:")
+    print(f"  {CLASS_NAMES[0]:>10} {CLASS_NAMES[1]:>10}")
+    for i, row in enumerate(metrics['confusion_matrix']):
+        print(f"  {CLASS_NAMES[i]:>10} {row[0]:>10} {row[1]:>10}")
+
+    print("\nClassification Report:")
+    print(classification_report(labels, preds, target_names=CLASS_NAMES))
+
+    return metrics
+
+
+def get_predictions_with_paths(
+    model: nn.Module,
+    dataset,
+    device: torch.device
+) -> list:
+    """Get predictions with image paths for error analysis."""
+    model.eval()
+    results = []
+
+    with torch.no_grad():
+        for idx in range(len(dataset)):
+            image, label = dataset[idx]
+            image = image.unsqueeze(0).to(device)
+
+            output = model(image)
+            prob = torch.sigmoid(output).item()
+            pred = 1 if prob > 0.5 else 0
+
+            results.append({
+                'path': dataset.image_paths[idx],
+                'true_label': label,
+                'pred_label': pred,
+                'probability': prob,
+                'correct': pred == label
+            })
+
+    return results

src/export.py

@@ -0,0 +1,190 @@
+"""
+ONNX export utilities for model deployment.
+
+ONNX (Open Neural Network Exchange) is a universal format that allows
+models to run on different frameworks and platforms:
+- TensorFlow, PyTorch, etc.
+- Mobile devices (iOS, Android)
+- Web browsers (ONNX.js)
+- C++, Java, and other languages
+- Optimized inference servers
+"""
+
+import torch
+import numpy as np
+from pathlib import Path
+from typing import Tuple, Optional
+
+from .config import CHECKPOINT_PATH, MODEL_DIR, IMAGE_SIZE
+from .model import create_model, get_device
+
+
+def export_to_onnx(
+    checkpoint_path: Path = CHECKPOINT_PATH,
+    output_path: Optional[Path] = None,
+    opset_version: int = 18
+) -> Path:
+    """
+    Export PyTorch model to ONNX format.
+
+    Args:
+        checkpoint_path: Path to the PyTorch checkpoint
+        output_path: Path for the ONNX model (default: models/best_model.onnx)
+        opset_version: ONNX opset version (14 is widely compatible)
+
+    Returns:
+        Path to the exported ONNX model
+    """
+    if output_path is None:
+        output_path = MODEL_DIR / "best_model.onnx"
+
+    # Load model
+    device = torch.device("cpu")  # Export on CPU for compatibility
+    model = create_model(pretrained=False, freeze_backbone=False, device=device)
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+
+    # Create dummy input (batch_size=1, channels=3, height=224, width=224)
+    dummy_input = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE)
+
+    # Export to ONNX
+    torch.onnx.export(
+        model,
+        dummy_input,
+        output_path,
+        export_params=True,
+        opset_version=opset_version,
+        do_constant_folding=True,  # Optimize constants
+        input_names=['image'],
+        output_names=['logits'],
+        dynamic_axes={
+            'image': {0: 'batch_size'},   # Variable batch size
+            'logits': {0: 'batch_size'}
+        }
+    )
+
+    print(f"Model exported to: {output_path}")
+    print(f"File size: {output_path.stat().st_size / 1024 / 1024:.2f} MB")
+
+    return output_path
+
+
+def validate_onnx_model(
+    onnx_path: Path,
+    checkpoint_path: Path = CHECKPOINT_PATH,
+    rtol: float = 1e-3,
+    atol: float = 1e-5
+) -> bool:
+    """
+    Validate that ONNX model produces same outputs as PyTorch model.
+
+    Args:
+        onnx_path: Path to ONNX model
+        checkpoint_path: Path to PyTorch checkpoint
+        rtol: Relative tolerance for comparison
+        atol: Absolute tolerance for comparison
+
+    Returns:
+        True if outputs match, False otherwise
+    """
+    import onnx
+    import onnxruntime as ort
+
+    # Check ONNX model is valid
+    onnx_model = onnx.load(onnx_path)
+    onnx.checker.check_model(onnx_model)
+    print("ONNX model structure is valid")
+
+    # Load PyTorch model
+    device = torch.device("cpu")
+    model = create_model(pretrained=False, freeze_backbone=False, device=device)
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+
+    # Create test input
+    test_input = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE)
+
+    # Get PyTorch output
+    with torch.no_grad():
+        pytorch_output = model(test_input).numpy()
+
+    # Get ONNX output
+    ort_session = ort.InferenceSession(str(onnx_path))
+    onnx_output = ort_session.run(
+        None,
+        {'image': test_input.numpy()}
+    )[0]
+
+    # Compare outputs
+    is_close = np.allclose(pytorch_output, onnx_output, rtol=rtol, atol=atol)
+
+    if is_close:
+        print("Validation PASSED: ONNX outputs match PyTorch outputs")
+        print(f"  PyTorch output: {pytorch_output.flatten()[:5]}...")
+        print(f"  ONNX output:    {onnx_output.flatten()[:5]}...")
+    else:
+        print("Validation FAILED: Outputs do not match!")
+        print(f"  Max difference: {np.max(np.abs(pytorch_output - onnx_output))}")
+
+    return is_close
+
+
+def predict_with_onnx(
+    onnx_path: Path,
+    image_tensor: np.ndarray
+) -> Tuple[str, float]:
+    """
+    Run inference using ONNX Runtime.
+
+    Args:
+        onnx_path: Path to ONNX model
+        image_tensor: Preprocessed image as numpy array (1, 3, 224, 224)
+
+    Returns:
+        Tuple of (predicted_class, confidence)
+    """
+    import onnxruntime as ort
+    from .config import CLASS_NAMES
+
+    # Create session
+    ort_session = ort.InferenceSession(str(onnx_path))
+
+    # Run inference
+    logits = ort_session.run(
+        None,
+        {'image': image_tensor.astype(np.float32)}
+    )[0]
+
+    # Apply sigmoid and get prediction
+    prob = 1 / (1 + np.exp(-logits[0, 0]))  # Sigmoid
+    pred_class = CLASS_NAMES[1] if prob > 0.5 else CLASS_NAMES[0]
+    confidence = float(prob if prob > 0.5 else 1 - prob)
+
+    return pred_class, confidence
+
+
+if __name__ == "__main__":
+    # Export model
+    print("=" * 50)
+    print("EXPORTING MODEL TO ONNX")
+    print("=" * 50)
+
+    onnx_path = export_to_onnx()
+
+    print("\n" + "=" * 50)
+    print("VALIDATING ONNX MODEL")
+    print("=" * 50)
+
+    validate_onnx_model(onnx_path)
+
+    print("\n" + "=" * 50)
+    print("TESTING ONNX INFERENCE")
+    print("=" * 50)
+
+    # Test with random input
+    test_input = np.random.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE).astype(np.float32)
+    pred_class, confidence = predict_with_onnx(onnx_path, test_input)
+    print(f"Test prediction: {pred_class} ({confidence:.1%})")

src/gradcam.py

@@ -0,0 +1,137 @@
+"""
+Grad-CAM visualization for model interpretability.
+"""
+
+import torch
+import numpy as np
+from PIL import Image
+from pathlib import Path
+from typing import Union
+import matplotlib.pyplot as plt
+
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+from .dataset import get_transforms
+from .config import IMAGENET_MEAN, IMAGENET_STD, CLASS_NAMES
+
+
+def get_gradcam(model, target_layer=None):
+    """Create GradCAM object for the model."""
+    if target_layer is None:
+        # Use the last conv layer of EfficientNet
+        target_layer = model.backbone.features[-1]
+    return GradCAM(model=model, target_layers=[target_layer])
+
+
+def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
+    """Denormalize tensor to numpy image [0,1]."""
+    mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
+    std = torch.tensor(IMAGENET_STD).view(3, 1, 1)
+    img = tensor.cpu() * std + mean
+    img = img.permute(1, 2, 0).numpy()
+    return np.clip(img, 0, 1)
+
+
+def generate_gradcam(
+    model,
+    image: Union[str, Path, Image.Image],
+    device: torch.device
+) -> tuple:
+    """Generate Grad-CAM heatmap for an image."""
+    model.eval()
+
+    # Load and transform image
+    if isinstance(image, (str, Path)):
+        image = Image.open(image).convert('RGB')
+
+    transform = get_transforms(is_training=False)
+    img_tensor = transform(image).unsqueeze(0).to(device)
+
+    # Get prediction
+    with torch.no_grad():
+        output = model(img_tensor)
+        prob = torch.sigmoid(output).item()
+
+    pred_class = CLASS_NAMES[1] if prob > 0.5 else CLASS_NAMES[0]
+    confidence = prob if prob > 0.5 else 1 - prob
+
+    # Generate Grad-CAM
+    cam = get_gradcam(model)
+    grayscale_cam = cam(input_tensor=img_tensor, targets=None)[0]
+
+    # Create visualization
+    rgb_img = denormalize_image(img_tensor[0])
+    cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
+
+    return cam_image, pred_class, confidence, rgb_img
+
+
+def plot_gradcam(
+    model,
+    image_path: Union[str, Path],
+    true_label: str,
+    device: torch.device,
+    save_path: str = None
+):
+    """Plot original image with Grad-CAM overlay."""
+    cam_image, pred_class, confidence, original = generate_gradcam(model, image_path, device)
+
+    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
+
+    # Original
+    axes[0].imshow(original)
+    axes[0].set_title(f"Original\nTrue: {true_label}")
+    axes[0].axis('off')
+
+    # Grad-CAM
+    color = 'green' if pred_class == true_label else 'red'
+    axes[1].imshow(cam_image)
+    axes[1].set_title(f"Grad-CAM\nPred: {pred_class} ({confidence:.1%})", color=color)
+    axes[1].axis('off')
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+
+    plt.show()
+    return pred_class, confidence
+
+
+def plot_gradcam_grid(
+    model,
+    image_paths: list,
+    true_labels: list,
+    device: torch.device,
+    save_path: str = None,
+    title: str = "Grad-CAM Visualizations"
+):
+    """Plot grid of Grad-CAM visualizations."""
+    n = len(image_paths)
+    fig, axes = plt.subplots(n, 2, figsize=(8, 3 * n))
+
+    if n == 1:
+        axes = axes.reshape(1, -1)
+
+    for i, (path, true_label) in enumerate(zip(image_paths, true_labels)):
+        cam_image, pred_class, confidence, original = generate_gradcam(model, path, device)
+
+        # Original
+        axes[i, 0].imshow(original)
+        axes[i, 0].set_title(f"True: {true_label}")
+        axes[i, 0].axis('off')
+
+        # Grad-CAM
+        color = 'green' if pred_class == true_label else 'red'
+        axes[i, 1].imshow(cam_image)
+        axes[i, 1].set_title(f"Pred: {pred_class} ({confidence:.1%})", color=color)
+        axes[i, 1].axis('off')
+
+    plt.suptitle(title, fontsize=14, fontweight='bold')
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+
+    plt.show()

src/model.py

@@ -0,0 +1,87 @@
+"""
+EfficientNet-B0 model for Pneumonia classification.
+"""
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from typing import Tuple
+
+from .config import DROPOUT_RATE, NUM_CLASSES
+
+
+class PneumoniaClassifier(nn.Module):
+    """EfficientNet-B0 based classifier for chest X-ray pneumonia detection."""
+
+    def __init__(
+        self,
+        pretrained: bool = True,
+        dropout_rate: float = DROPOUT_RATE,
+        freeze_backbone: bool = True
+    ):
+        super().__init__()
+
+        # Load pretrained EfficientNet-B0
+        weights = models.EfficientNet_B0_Weights.IMAGENET1K_V1 if pretrained else None
+        self.backbone = models.efficientnet_b0(weights=weights)
+
+        # Get the number of features from the classifier
+        in_features = self.backbone.classifier[1].in_features  # 1280
+
+        # Replace classifier head
+        self.backbone.classifier = nn.Sequential(
+            nn.Dropout(p=dropout_rate, inplace=True),
+            nn.Linear(in_features, NUM_CLASSES)
+        )
+
+        # Freeze backbone if specified
+        if freeze_backbone:
+            self.freeze_backbone()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.backbone(x)
+
+    def freeze_backbone(self):
+        """Freeze all layers except the classifier."""
+        for param in self.backbone.features.parameters():
+            param.requires_grad = False
+
+    def unfreeze_backbone(self):
+        """Unfreeze all layers for fine-tuning."""
+        for param in self.backbone.features.parameters():
+            param.requires_grad = True
+
+    def get_param_counts(self) -> Tuple[int, int]:
+        """Return (trainable_params, total_params)."""
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in self.parameters())
+        return trainable, total
+
+
+def create_model(
+    pretrained: bool = True,
+    dropout_rate: float = DROPOUT_RATE,
+    freeze_backbone: bool = True,
+    device: str = None
+) -> PneumoniaClassifier:
+    """Factory function to create the model."""
+    if device is None:
+        device = "mps" if torch.backends.mps.is_available() else \
+                 "cuda" if torch.cuda.is_available() else "cpu"
+
+    model = PneumoniaClassifier(
+        pretrained=pretrained,
+        dropout_rate=dropout_rate,
+        freeze_backbone=freeze_backbone
+    )
+
+    return model.to(device)
+
+
+def get_device() -> torch.device:
+    """Get the best available device."""
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    elif torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")

src/predict.py

@@ -0,0 +1,47 @@
+"""
+Inference functions for Pneumonia classification.
+"""
+
+import torch
+import torch.nn as nn
+from PIL import Image
+from pathlib import Path
+from typing import Union, Tuple
+
+from .dataset import get_transforms
+from .config import CLASS_NAMES, CHECKPOINT_PATH
+
+
+def load_model(model: nn.Module, checkpoint_path: Path = CHECKPOINT_PATH, device: str = "cpu") -> nn.Module:
+    """Load model from checkpoint."""
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    return model
+
+
+def predict_image(
+    model: nn.Module,
+    image: Union[str, Path, Image.Image],
+    device: torch.device
+) -> Tuple[str, float]:
+    """Predict class for a single image."""
+    model.eval()
+
+    # Load image if path
+    if isinstance(image, (str, Path)):
+        image = Image.open(image).convert('RGB')
+
+    # Transform
+    transform = get_transforms(is_training=False)
+    img_tensor = transform(image).unsqueeze(0).to(device)
+
+    # Predict
+    with torch.no_grad():
+        output = model(img_tensor)
+        prob = torch.sigmoid(output).item()
+
+    pred_class = CLASS_NAMES[1] if prob > 0.5 else CLASS_NAMES[0]
+    confidence = prob if prob > 0.5 else 1 - prob
+
+    return pred_class, confidence

src/train.py

@@ -0,0 +1,250 @@
+"""
+Training pipeline for Pneumonia classification.
+"""
+
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torch.utils.data import DataLoader
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+import time
+
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
+
+from .config import (
+    STAGE1_EPOCHS, STAGE1_LR, STAGE2_EPOCHS, STAGE2_LR,
+    WEIGHT_DECAY, SCHEDULER_PATIENCE, SCHEDULER_FACTOR,
+    EARLY_STOP_PATIENCE, CHECKPOINT_PATH, MODEL_DIR
+)
+from .model import PneumoniaClassifier, get_device
+
+
+class EarlyStopping:
+    """Early stopping to prevent overfitting."""
+
+    def __init__(self, patience: int = 7, min_delta: float = 0.001):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.counter = 0
+        self.best_loss = float('inf')
+        self.should_stop = False
+
+    def __call__(self, val_loss: float) -> bool:
+        if val_loss < self.best_loss - self.min_delta:
+            self.best_loss = val_loss
+            self.counter = 0
+        else:
+            self.counter += 1
+            if self.counter >= self.patience:
+                self.should_stop = True
+        return self.should_stop
+
+
+def train_epoch(
+    model: nn.Module,
+    loader: DataLoader,
+    criterion: nn.Module,
+    optimizer: torch.optim.Optimizer,
+    device: torch.device
+) -> Tuple[float, float]:
+    """Train for one epoch."""
+    model.train()
+    total_loss = 0
+    all_preds, all_labels = [], []
+
+    for images, labels in loader:
+        images = images.to(device)
+        labels = labels.float().unsqueeze(1).to(device)
+
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss.item() * images.size(0)
+        preds = (torch.sigmoid(outputs) > 0.5).int()
+        all_preds.extend(preds.cpu().numpy())
+        all_labels.extend(labels.cpu().numpy())
+
+    avg_loss = total_loss / len(loader.dataset)
+    accuracy = accuracy_score(all_labels, all_preds)
+    return avg_loss, accuracy
+
+
+def validate(
+    model: nn.Module,
+    loader: DataLoader,
+    criterion: nn.Module,
+    device: torch.device
+) -> Dict[str, float]:
+    """Validate the model."""
+    model.eval()
+    total_loss = 0
+    all_preds, all_labels = [], []
+
+    with torch.no_grad():
+        for images, labels in loader:
+            images = images.to(device)
+            labels = labels.float().unsqueeze(1).to(device)
+
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            total_loss += loss.item() * images.size(0)
+            preds = (torch.sigmoid(outputs) > 0.5).int()
+            all_preds.extend(preds.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+
+    avg_loss = total_loss / len(loader.dataset)
+
+    return {
+        'loss': avg_loss,
+        'accuracy': accuracy_score(all_labels, all_preds),
+        'precision': precision_score(all_labels, all_preds, zero_division=0),
+        'recall': recall_score(all_labels, all_preds, zero_division=0),
+        'f1': f1_score(all_labels, all_preds, zero_division=0)
+    }
+
+
+def train(
+    model: PneumoniaClassifier,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    pos_weight: torch.Tensor,
+    epochs: int,
+    lr: float,
+    device: torch.device,
+    stage: str = "stage1",
+    use_wandb: bool = True,
+    wandb_run = None
+) -> Dict[str, list]:
+    """Training loop with validation."""
+
+    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight.to(device))
+    optimizer = AdamW(
+        filter(lambda p: p.requires_grad, model.parameters()),
+        lr=lr,
+        weight_decay=WEIGHT_DECAY
+    )
+    scheduler = ReduceLROnPlateau(
+        optimizer, mode='min',
+        patience=SCHEDULER_PATIENCE,
+        factor=SCHEDULER_FACTOR
+    )
+    early_stopping = EarlyStopping(patience=EARLY_STOP_PATIENCE)
+
+    history = {'train_loss': [], 'val_loss': [], 'val_acc': [], 'val_f1': [], 'lr': []}
+    best_val_loss = float('inf')
+
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+
+    for epoch in range(epochs):
+        start = time.time()
+
+        # Train
+        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
+
+        # Validate
+        val_metrics = validate(model, val_loader, criterion, device)
+
+        # Get current LR
+        current_lr = optimizer.param_groups[0]['lr']
+
+        # Update scheduler
+        scheduler.step(val_metrics['loss'])
+
+        # Log
+        elapsed = time.time() - start
+        print(f"[{stage}] Epoch {epoch+1}/{epochs} ({elapsed:.1f}s) | "
+              f"Train Loss: {train_loss:.4f} | "
+              f"Val Loss: {val_metrics['loss']:.4f} | "
+              f"Val Acc: {val_metrics['accuracy']:.3f} | "
+              f"Val F1: {val_metrics['f1']:.3f} | "
+              f"LR: {current_lr:.2e}")
+
+        # W&B logging
+        if use_wandb and wandb_run:
+            wandb_run.log({
+                f"{stage}/train_loss": train_loss,
+                f"{stage}/train_acc": train_acc,
+                f"{stage}/val_loss": val_metrics['loss'],
+                f"{stage}/val_acc": val_metrics['accuracy'],
+                f"{stage}/val_precision": val_metrics['precision'],
+                f"{stage}/val_recall": val_metrics['recall'],
+                f"{stage}/val_f1": val_metrics['f1'],
+                f"{stage}/lr": current_lr,
+                "epoch": epoch + 1
+            })
+
+        # Save history
+        history['train_loss'].append(train_loss)
+        history['val_loss'].append(val_metrics['loss'])
+        history['val_acc'].append(val_metrics['accuracy'])
+        history['val_f1'].append(val_metrics['f1'])
+        history['lr'].append(current_lr)
+
+        # Save best model
+        if val_metrics['loss'] < best_val_loss:
+            best_val_loss = val_metrics['loss']
+            torch.save({
+                'epoch': epoch + 1,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_loss': best_val_loss,
+                'val_metrics': val_metrics
+            }, CHECKPOINT_PATH)
+            print(f"  -> Saved best model (val_loss: {best_val_loss:.4f})")
+
+        # Early stopping
+        if early_stopping(val_metrics['loss']):
+            print(f"Early stopping triggered at epoch {epoch+1}")
+            break
+
+    return history
+
+
+def train_two_stage(
+    model: PneumoniaClassifier,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    pos_weight: torch.Tensor,
+    device: torch.device,
+    use_wandb: bool = True,
+    wandb_run = None
+) -> Dict[str, list]:
+    """Two-stage training: frozen backbone then fine-tuning."""
+
+    # Stage 1: Train classifier only
+    print("\n" + "=" * 60)
+    print("STAGE 1: Training classifier (backbone frozen)")
+    print("=" * 60)
+    model.freeze_backbone()
+    trainable, total = model.get_param_counts()
+    print(f"Trainable params: {trainable:,} / {total:,}")
+
+    history1 = train(
+        model, train_loader, val_loader, pos_weight,
+        epochs=STAGE1_EPOCHS, lr=STAGE1_LR, device=device,
+        stage="stage1", use_wandb=use_wandb, wandb_run=wandb_run
+    )
+
+    # Stage 2: Fine-tune entire network
+    print("\n" + "=" * 60)
+    print("STAGE 2: Fine-tuning entire network")
+    print("=" * 60)
+    model.unfreeze_backbone()
+    trainable, total = model.get_param_counts()
+    print(f"Trainable params: {trainable:,} / {total:,}")
+
+    history2 = train(
+        model, train_loader, val_loader, pos_weight,
+        epochs=STAGE2_EPOCHS, lr=STAGE2_LR, device=device,
+        stage="stage2", use_wandb=use_wandb, wandb_run=wandb_run
+    )
+
+    # Combine histories
+    history = {k: history1[k] + history2[k] for k in history1}
+    return history

src/utils.py

@@ -0,0 +1,74 @@
+"""
+Utility functions for visualization and helpers.
+"""
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from typing import Optional
+
+from .config import IMAGENET_MEAN, IMAGENET_STD, CLASS_NAMES
+
+
+def denormalize(tensor: torch.Tensor) -> torch.Tensor:
+    """Denormalize image tensor from ImageNet normalization."""
+    mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
+    std = torch.tensor(IMAGENET_STD).view(3, 1, 1)
+    return tensor * std + mean
+
+
+def show_batch(
+    images: torch.Tensor,
+    labels: torch.Tensor,
+    predictions: Optional[torch.Tensor] = None,
+    n_images: int = 8,
+    save_path: Optional[str] = None
+):
+    """Display a batch of images with labels."""
+    n_images = min(n_images, len(images))
+    cols = 4
+    rows = (n_images + cols - 1) // cols
+
+    fig, axes = plt.subplots(rows, cols, figsize=(12, 3 * rows))
+    axes = axes.flatten() if rows > 1 else [axes] if cols == 1 else axes
+
+    for idx in range(n_images):
+        img = denormalize(images[idx]).permute(1, 2, 0).numpy()
+        img = np.clip(img, 0, 1)
+
+        axes[idx].imshow(img)
+        axes[idx].axis('off')
+
+        label = CLASS_NAMES[labels[idx]]
+        title = f"True: {label}"
+
+        if predictions is not None:
+            pred = CLASS_NAMES[predictions[idx]]
+            color = 'green' if pred == label else 'red'
+            title += f"\nPred: {pred}"
+            axes[idx].set_title(title, color=color, fontsize=10)
+        else:
+            axes[idx].set_title(title, fontsize=10)
+
+    # Hide empty subplots
+    for idx in range(n_images, len(axes)):
+        axes[idx].axis('off')
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=150, bbox_inches='tight')
+
+    plt.show()
+
+
+def set_seed(seed: int = 42):
+    """Set random seed for reproducibility."""
+    import random
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    if torch.backends.mps.is_available():
+        torch.mps.manual_seed(seed)