Auto-deploy from GitHub: 495db78a06be79166200269bb14d9e9b1e8906d6

app.py

@@ -1,28 +1,24 @@
 """
-Hugging Face Spaces - Pneumonia Detection App
-This is a self-contained version for Hugging Face Spaces deployment.
-Copy this file as 'app.py' to your HF Spaces repository.
+Streamlit Web UI for Pneumonia Detection.
+
+Run with: streamlit run app/app.py
 """
 
+import sys
+from pathlib import Path
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
 import streamlit as st
 import torch
-import torch.nn as nn
-from torchvision import models, transforms
 from PIL import Image
-import numpy as np
-from pytorch_grad_cam import GradCAM
-from pytorch_grad_cam.utils.image import show_cam_on_image
 import time
 
-# =============================================================================
-# Configuration
-# =============================================================================
-
-IMAGE_SIZE = 224
-IMAGENET_MEAN = [0.485, 0.456, 0.406]
-IMAGENET_STD = [0.229, 0.224, 0.225]
-CLASS_NAMES = ["NORMAL", "PNEUMONIA"]
-MODEL_PATH = "models/best_model.pt"
+from src.config import CHECKPOINT_PATH, CLASS_NAMES
+from src.model import create_model, get_device
+from src.predict import load_model, predict_image
+from src.gradcam import generate_gradcam
 
 # =============================================================================
 # Page Configuration
@@ -31,79 +27,10 @@ MODEL_PATH = "models/best_model.pt"
 st.set_page_config(
     page_title="Pneumonia Detection",
     page_icon="🫁",
-    layout="wide"
+    layout="wide",
+    initial_sidebar_state="expanded"
 )
 
-# =============================================================================
-# Model Definition
-# =============================================================================
-
-class PneumoniaClassifier(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.backbone = models.efficientnet_b0(weights=None)
-        in_features = self.backbone.classifier[1].in_features
-        self.backbone.classifier = nn.Sequential(
-            nn.Dropout(p=0.3, inplace=True),
-            nn.Linear(in_features, 1)
-        )
-
-    def forward(self, x):
-        return self.backbone(x)
-
-
-# =============================================================================
-# Helper Functions
-# =============================================================================
-
-def get_transforms():
-    return transforms.Compose([
-        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
-    ])
-
-
-def denormalize(tensor):
-    mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
-    std = torch.tensor(IMAGENET_STD).view(3, 1, 1)
-    img = tensor.cpu() * std + mean
-    return img.permute(1, 2, 0).numpy().clip(0, 1)
-
-
-@st.cache_resource
-def load_model():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = PneumoniaClassifier()
-    checkpoint = torch.load(MODEL_PATH, map_location=device)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model.eval()
-    return model.to(device), device
-
-
-def predict(model, image, device):
-    transform = get_transforms()
-    img_tensor = transform(image).unsqueeze(0).to(device)
-
-    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, img_tensor
-
-
-def generate_gradcam(model, img_tensor, device):
-    target_layer = model.backbone.features[-1]
-    cam = GradCAM(model=model, target_layers=[target_layer])
-    grayscale_cam = cam(input_tensor=img_tensor, targets=None)[0]
-
-    rgb_img = denormalize(img_tensor[0])
-    cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
-    return cam_image, rgb_img
-
-
 # =============================================================================
 # Custom CSS
 # =============================================================================
@@ -137,17 +64,42 @@ st.markdown("""
         background-color: #FFEBEE;
         border: 2px solid #F44336;
     }
+    .confidence-text {
+        font-size: 1.2rem;
+        font-weight: bold;
+    }
+    .metric-card {
+        background-color: #f8f9fa;
+        padding: 1rem;
+        border-radius: 8px;
+        text-align: center;
+    }
 </style>
 """, unsafe_allow_html=True)
 
+# =============================================================================
+# Model Loading (Cached)
+# =============================================================================
+
+@st.cache_resource
+def load_model_cached():
+    """Load model once and cache it."""
+    device = get_device()
+    model = create_model(pretrained=False, freeze_backbone=False, device=device)
+    model = load_model(model, CHECKPOINT_PATH, device)
+    return model, device
+
+
 # =============================================================================
 # Sidebar
 # =============================================================================
 
 with st.sidebar:
-    st.title("🫁 About")
+    st.image("https://img.icons8.com/fluency/96/lungs.png", width=80)
+    st.title("About")
+
     st.markdown("""
-    Deep learning model for detecting **pneumonia** from chest X-rays.
+    This application uses deep learning to detect **pneumonia** from chest X-ray images.
 
     **Model:** EfficientNet-B0
     **Accuracy:** 90.5%
@@ -155,101 +107,172 @@ with st.sidebar:
     """)
 
     st.divider()
+
+    st.subheader("How to Use")
+    st.markdown("""
+    1. Upload a chest X-ray image
+    2. Click **Analyze Image**
+    3. View prediction and Grad-CAM
+    """)
+
+    st.divider()
+
     st.subheader("Model Metrics")
     col1, col2 = st.columns(2)
-    col1.metric("Accuracy", "90.5%")
-    col2.metric("Recall", "98.2%")
+    with col1:
+        st.metric("Accuracy", "90.5%")
+        st.metric("Precision", "88.0%")
+    with col2:
+        st.metric("Recall", "98.2%")
+        st.metric("F1 Score", "92.8%")
 
     st.divider()
-    st.markdown("*Built with PyTorch & Streamlit*")
+
+    st.markdown("""
+    **Links:**
+    [GitHub Repository](#) | [Live Demo](#)
+
+    ---
+    *Built with PyTorch & Streamlit*
+    """)
 
 # =============================================================================
 # Main Content
 # =============================================================================
 
-st.markdown('<p class="main-header">🫁 Pneumonia Detection</p>', unsafe_allow_html=True)
-st.markdown('<p class="sub-header">Upload a chest X-ray image to detect pneumonia</p>', unsafe_allow_html=True)
+# Header
+st.markdown('<p class="main-header">🫁 Pneumonia Detection from Chest X-Rays</p>', unsafe_allow_html=True)
+st.markdown('<p class="sub-header">Upload a chest X-ray image to detect pneumonia using AI</p>', unsafe_allow_html=True)
 
 # Load model
 try:
-    model, device = load_model()
+    model, device = load_model_cached()
     model_loaded = True
 except Exception as e:
     st.error(f"Failed to load model: {e}")
     model_loaded = False
 
 if model_loaded:
+    # Create columns for layout
     col1, col2 = st.columns([1, 1])
 
     with col1:
         st.subheader("📤 Upload Image")
+
         uploaded_file = st.file_uploader(
             "Choose a chest X-ray image",
-            type=["jpg", "jpeg", "png"]
+            type=["jpg", "jpeg", "png"],
+            help="Supported formats: JPG, JPEG, PNG"
         )
 
-        # Sample images
-        st.subheader("🖼️ Or Try Sample Images")
+        # Sample images section
+        st.markdown("---")
+        st.markdown("**Or try a sample image:**")
+
         sample_col1, sample_col2 = st.columns(2)
+
+        use_sample = None
         with sample_col1:
-            if st.button("Normal Sample", use_container_width=True):
-                st.session_state.sample = "samples/normal_sample.jpeg"
+            if st.button("🟢 Normal Sample", width="stretch"):
+                use_sample = "normal"
         with sample_col2:
-            if st.button("Pneumonia Sample", use_container_width=True):
-                st.session_state.sample = "samples/pneumonia_sample.jpeg"
-
-    # Determine which image to use
-    image = None
+            if st.button("🔴 Pneumonia Sample", width="stretch"):
+                use_sample = "pneumonia"
+
+        # Load sample image if selected
+        if use_sample == "normal":
+            sample_path = Path(__file__).parent / "samples" / "normal_sample.jpeg"
+            if sample_path.exists():
+                uploaded_file = sample_path
+        elif use_sample == "pneumonia":
+            sample_path = Path(__file__).parent / "samples" / "pneumonia_sample.jpeg"
+            if sample_path.exists():
+                uploaded_file = sample_path
+
+    with col2:
+        st.subheader("🔍 Analysis Results")
+        results_placeholder = st.empty()
+
+    # Process image if uploaded
     if uploaded_file is not None:
-        image = Image.open(uploaded_file).convert("RGB")
-        st.session_state.sample = None  # Clear sample when uploading
-    elif "sample" in st.session_state and st.session_state.sample:
-        image = Image.open(st.session_state.sample).convert("RGB")
-
-    if image is not None:
+        # Load image
+        if isinstance(uploaded_file, Path):
+            image = Image.open(uploaded_file).convert("RGB")
+            st.session_state['image_source'] = str(uploaded_file)
+        else:
+            image = Image.open(uploaded_file).convert("RGB")
+            st.session_state['image_source'] = uploaded_file.name
+
+        # Display uploaded image
+        with col1:
+            st.image(image, caption="Uploaded X-Ray", width="stretch")
 
+        # Analyze button
         with col1:
-            st.image(image, caption="Uploaded X-Ray", use_container_width=True)
-            analyze = st.button("🔬 Analyze Image", type="primary", use_container_width=True)
+            analyze_button = st.button("🔬 Analyze Image", type="primary", width="stretch")
 
-        if analyze:
+        if analyze_button:
             with col2:
-                with st.spinner("Analyzing..."):
+                with st.spinner("Analyzing image..."):
+                    # Run prediction
                     start_time = time.time()
-                    pred_class, confidence, img_tensor = predict(model, image, device)
-                    cam_image, original = generate_gradcam(model, img_tensor, device)
+                    pred_class, confidence = predict_image(model, image, device)
                     inference_time = (time.time() - start_time) * 1000
 
-                # Results
+                    # Generate Grad-CAM
+                    cam_image, _, _, original = generate_gradcam(model, image, device)
+
+                # Display results
                 if pred_class == "PNEUMONIA":
                     st.markdown(f"""
                     <div class="prediction-box prediction-pneumonia">
-                        <h2 style="color: #F44336;">⚠️ PNEUMONIA DETECTED</h2>
-                        <p>Confidence: {confidence:.1%}</p>
+                        <h2 style="color: #F44336; margin: 0;">⚠️ PNEUMONIA DETECTED</h2>
+                        <p class="confidence-text">Confidence: {confidence:.1%}</p>
                     </div>
                     """, unsafe_allow_html=True)
                 else:
                     st.markdown(f"""
                     <div class="prediction-box prediction-normal">
-                        <h2 style="color: #4CAF50;">✅ NORMAL</h2>
-                        <p>Confidence: {confidence:.1%}</p>
+                        <h2 style="color: #4CAF50; margin: 0;">✅ NORMAL</h2>
+                        <p class="confidence-text">Confidence: {confidence:.1%}</p>
                     </div>
                     """, unsafe_allow_html=True)
 
                 # Metrics row
                 m1, m2, m3 = st.columns(3)
-                m1.metric("Prediction", pred_class)
-                m2.metric("Confidence", f"{confidence:.1%}")
-                m3.metric("Time", f"{inference_time:.0f}ms")
+                with m1:
+                    st.metric("Prediction", pred_class)
+                with m2:
+                    st.metric("Confidence", f"{confidence:.1%}")
+                with m3:
+                    st.metric("Time", f"{inference_time:.0f}ms")
+
+                # Grad-CAM visualization
+                st.markdown("---")
+                st.subheader("🔥 Grad-CAM Visualization")
+                st.caption("Highlighted regions show areas that influenced the prediction")
 
-                # Grad-CAM
-                st.subheader("🔥 Grad-CAM")
                 gcol1, gcol2 = st.columns(2)
-                gcol1.image(original, caption="Original", use_container_width=True)
-                gcol2.image(cam_image, caption="Heatmap", use_container_width=True)
+                with gcol1:
+                    st.image(original, caption="Original", width="stretch")
+                with gcol2:
+                    st.image(cam_image, caption="Grad-CAM Heatmap", width="stretch")
+
+                # Disclaimer
+                st.warning("""
+                **Disclaimer:** This tool is for educational purposes only and should not be used
+                for medical diagnosis. Always consult a qualified healthcare professional.
+                """)
 
-                st.warning("**Disclaimer:** For educational purposes only. Consult a healthcare professional.")
+else:
+    st.error("Model could not be loaded. Please check the model file exists.")
+
+# =============================================================================
+# Footer
+# =============================================================================
 
-# Footers
 st.markdown("---")
-st.markdown("<p style='text-align:center;color:#888;'>Built with PyTorch & Streamlit</p>", unsafe_allow_html=True)
+st.markdown(
+    "<p style='text-align: center; color: #888;'>Built with ❤️ using PyTorch, EfficientNet-B0, and Streamlit</p>",
+    unsafe_allow_html=True
+)

requirements.txt

@@ -1,20 +1,45 @@
-# Hugging Face Spaces Requirements
-# Minimal dependencies for deployment
-
-# PyTorch (CPU version for HF Spaces)
---extra-index-url https://download.pytorch.org/whl/cpu
-torch>=2.5.0
-torchvision>=0.20.0
-
-# Core
+# Core Deep Learning
+torch>=2.1.0
+torchvision>=0.16.0
 pillow>=10.0.0
 numpy>=1.24.0
 
-# OpenCV headless (must be before grad-cam to avoid libGL issues)
-opencv-python-headless>=4.8.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
+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)