{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0",
   "mimetype": "text/x-python",
   "file_extension": ".py",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python"
  },
  "colab": {
   "provenance": [],
   "gpuType": "T4",
   "name": "DermaScan AI: Vision Encoders for Dermatology Screening -- Implementation Notebook"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 🔧 Setup: Run this cell first!\n",
    "# Check GPU availability and install dependencies\n",
    "\n",
    "import torch\n",
    "import sys\n",
    "\n",
    "# Check GPU\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "    print(f\"✅ GPU available: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "    print(\"⚠️ No GPU detected. Some cells may run slowly.\")\n",
    "    print(\"   Go to Runtime → Change runtime type → GPU\")\n",
    "\n",
    "print(f\"\\n📦 Python {sys.version.split()[0]}\")\n",
    "print(f\"🔥 PyTorch {torch.__version__}\")\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "SEED = 42\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(SEED)\n",
    "\n",
    "print(f\"🎲 Random seed set to {SEED}\")\n",
    "\n",
    "%matplotlib inline"
   ],
   "id": "setup_cell"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "<cell_type>markdown</cell_type># DermaScan AI: Vision Encoders for Automated Dermatology Screening\n\n**Implementation Notebook -- Vizuara Case Study**\n\nIn this notebook, we implement the full pipeline for DermaScan AI's automated dermatology screening system using real dermatoscopic images from DermaMNIST:\n1. Data loading with class-weighted sampling\n2. CNN (ResNet-18) and ViT (ViT-B/16) encoders\n3. Custom loss function (weighted CE + focal loss)\n4. Training, evaluation, and error analysis\n5. Deployment optimization\n\n**Dataset:** DermaMNIST contains 10,015 real dermatoscopic images across 7 skin lesion types, sourced from the HAM10000 dataset used in clinical research. This is the real version of the task DermaScan AI needs to solve.\n\n**Runtime:** Google Colab (T4 GPU required)\n**Estimated time:** 90-120 minutes"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# Setup and dependencies\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nimport torchvision.transforms as transforms\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import confusion_matrix, classification_report, roc_curve, roc_auc_score\nfrom torch.utils.data import Dataset\nimport time\n\n!pip install -q medmnist\nimport medmnist\nfrom medmnist import DermaMNIST\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f'Using device: {device}')\n\n# DermaMNIST: real dermatoscopic images, 7 classes\n# Classes: 0=actinic_keratosis, 1=basal_cell_carcinoma, 2=benign_keratosis,\n#          3=dermatofibroma, 4=melanoma, 5=melanocytic_nevus, 6=vascular_lesion\nNUM_CLASSES = 7\nCLASS_NAMES = ['actinic_keratosis', 'basal_cell_carcinoma', 'benign_keratosis',\n               'dermatofibroma', 'melanoma', 'melanocytic_nevus', 'vascular_lesion']\nMELANOMA_IDX = 4  # Melanoma is class 4 in DermaMNIST"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "<cell_type>markdown</cell_type>## Section 3.1: Data Loading and Preprocessing\n\n**TODO 1:** Implement the data loading pipeline with class-weighted sampling to handle the severe class imbalance in DermaMNIST (melanocytic nevi dominate the dataset)."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "class IntLabelDataset(Dataset):\n    \"\"\"Wraps a MedMNIST dataset to return integer labels and support .targets attribute.\"\"\"\n    def __init__(self, dataset):\n        self.dataset = dataset\n        # Pre-extract all labels for WeightedRandomSampler\n        self.targets = [int(dataset[i][1].item()) for i in range(len(dataset))]\n    def __len__(self):\n        return len(self.dataset)\n    def __getitem__(self, idx):\n        img, label = self.dataset[idx]\n        return img, int(label.item())\n\ndef load_dermascan_dataset(img_size=224):\n    \"\"\"\n    Load and preprocess the DermaMNIST dermatoscopic image dataset.\n\n    Returns:\n        train_loader, val_loader, test_loader, class_weights\n    \"\"\"\n    # Training transforms with medical imaging augmentation\n    transform_train = transforms.Compose([\n        transforms.Resize((img_size, img_size)),\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomVerticalFlip(),\n        transforms.RandomRotation(30),\n        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1),\n        transforms.ToTensor(),\n        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n        transforms.RandomErasing(p=0.1),\n    ])\n\n    transform_test = transforms.Compose([\n        transforms.Resize((img_size, img_size)),\n        transforms.ToTensor(),\n        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n    ])\n\n    # Load real DermaMNIST data\n    train_raw = DermaMNIST(split='train', download=True, transform=transform_train)\n    test_raw = DermaMNIST(split='test', download=True, transform=transform_test)\n\n    trainset = IntLabelDataset(train_raw)\n    testset = IntLabelDataset(test_raw)\n\n    # Compute class weights\n    class_counts = np.bincount(trainset.targets, minlength=NUM_CLASSES)\n    class_weights = len(trainset.targets) / (NUM_CLASSES * class_counts + 1e-6)\n    class_weights = torch.FloatTensor(class_weights)\n    print(f'Class distribution: {dict(zip(CLASS_NAMES, class_counts))}')\n    print(f'Class weights: {class_weights.numpy().round(2)}')\n\n    # Weighted sampler for balanced batches\n    sample_weights = class_weights[trainset.targets]\n    sampler = torch.utils.data.WeightedRandomSampler(\n        weights=sample_weights, num_samples=len(trainset), replacement=True)\n\n    train_loader = torch.utils.data.DataLoader(\n        trainset, batch_size=64, sampler=sampler, num_workers=2)\n\n    # Split test into val/test\n    val_size = len(testset) // 2\n    val_set, test_set = torch.utils.data.random_split(\n        testset, [val_size, len(testset) - val_size])\n    val_loader = torch.utils.data.DataLoader(val_set, batch_size=64, num_workers=2)\n    test_loader = torch.utils.data.DataLoader(test_set, batch_size=64, num_workers=2)\n\n    return train_loader, val_loader, test_loader, class_weights\n\ntrain_loader, val_loader, test_loader, class_weights = load_dermascan_dataset(img_size=32)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "<cell_type>markdown</cell_type>## Section 3.2: Exploratory Data Analysis\n\n**TODO 2:** Visualize the DermaMNIST dataset to understand class distribution and image characteristics."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_eda(train_loader, class_names):\n",
    "    \"\"\"\n",
    "    TODO: Implement exploratory data analysis.\n",
    "    \n",
    "    1. Plot class distribution bar chart\n",
    "    2. Display sample images from each class\n",
    "    3. Compute mean pixel intensity per class\n",
    "    \"\"\"\n",
    "    # Collect a batch of images\n",
    "    images, labels = next(iter(train_loader))\n",
    "    \n",
    "    # Plot class distribution\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(14, 4))\n",
    "    \n",
    "    unique, counts = np.unique(labels.numpy(), return_counts=True)\n",
    "    axes[0].bar(range(len(class_names)), [counts[counts==i].sum() if i in unique else 0 \n",
    "                for i in range(len(class_names))], color='steelblue')\n",
    "    axes[0].set_xticks(range(len(class_names)))\n",
    "    axes[0].set_xticklabels(class_names, rotation=45, ha='right', fontsize=8)\n",
    "    axes[0].set_ylabel('Count (batch)')\n",
    "    axes[0].set_title('Class Distribution in Batch')\n",
    "    \n",
    "    # Show sample images\n",
    "    fig2, axes2 = plt.subplots(2, 7, figsize=(16, 5))\n",
    "    fig2.suptitle('Sample Images per Class', fontsize=14)\n",
    "    for cls in range(min(7, len(class_names))):\n",
    "        mask = labels == cls\n",
    "        if mask.sum() > 0:\n",
    "            idx = mask.nonzero()[0][0]\n",
    "            img = images[idx].permute(1, 2, 0).numpy()\n",
    "            img = (img - img.min()) / (img.max() - img.min())\n",
    "            axes2[0, cls].imshow(img)\n",
    "            axes2[0, cls].set_title(class_names[cls], fontsize=8)\n",
    "            axes2[0, cls].axis('off')\n",
    "            if mask.sum() > 1:\n",
    "                idx2 = mask.nonzero()[0][1]\n",
    "                img2 = images[idx2].permute(1, 2, 0).numpy()\n",
    "                img2 = (img2 - img2.min()) / (img2.max() - img2.min())\n",
    "                axes2[1, cls].imshow(img2)\n",
    "            axes2[1, cls].axis('off')\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "run_eda(train_loader, CLASS_NAMES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3.3: Baseline Model (ResNet-50)\n",
    "\n",
    "**TODO 3:** Create the ResNet-50 baseline with pretrained ImageNet weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_baseline_model(num_classes=7, pretrained=True, img_size=32):\n",
    "    \"\"\"\n",
    "    TODO: Create a ResNet-50 baseline model.\n",
    "    \n",
    "    Steps:\n",
    "    1. Load ResNet-50 with pretrained weights\n",
    "    2. Replace final fc layer for num_classes outputs\n",
    "    3. If img_size < 224, modify first conv to handle smaller inputs\n",
    "    \"\"\"\n",
    "    # For small images (CIFAR), use ResNet-18 instead\n",
    "    if img_size <= 64:\n",
    "        model = torchvision.models.resnet18(weights='IMAGENET1K_V1' if pretrained else None)\n",
    "        model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        model.maxpool = nn.Identity()\n",
    "        model.fc = nn.Linear(512, num_classes)\n",
    "    else:\n",
    "        model = torchvision.models.resnet50(weights='IMAGENET1K_V2' if pretrained else None)\n",
    "        model.fc = nn.Linear(2048, num_classes)\n",
    "    return model\n",
    "\n",
    "baseline = create_baseline_model(img_size=32).to(device)\n",
    "print(f'Baseline params: {sum(p.numel() for p in baseline.parameters()):,}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3.4: Vision Transformer Model\n",
    "\n",
    "**TODO 4:** Create the ViT model and implement the custom loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleViTForMedical(nn.Module):\n",
    "    \"\"\"Vision Transformer adapted for small medical images.\"\"\"\n",
    "    def __init__(self, img_size=32, patch_size=4, embed_dim=192,\n",
    "                 num_heads=6, num_layers=6, num_classes=7):\n",
    "        super().__init__()\n",
    "        self.num_patches = (img_size // patch_size) ** 2\n",
    "        self.patch_embed = nn.Conv2d(3, embed_dim, patch_size, stride=patch_size)\n",
    "        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n",
    "        self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches + 1, embed_dim) * 0.02)\n",
    "        encoder_layer = nn.TransformerEncoderLayer(\n",
    "            d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim * 4,\n",
    "            activation='gelu', batch_first=True, norm_first=True, dropout=0.1)\n",
    "        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)\n",
    "        self.norm = nn.LayerNorm(embed_dim)\n",
    "        self.head = nn.Sequential(\n",
    "            nn.Linear(embed_dim, embed_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(0.1),\n",
    "            nn.Linear(embed_dim, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        B = x.shape[0]\n",
    "        x = self.patch_embed(x).flatten(2).transpose(1, 2)\n",
    "        x = torch.cat([self.cls_token.expand(B, -1, -1), x], dim=1) + self.pos_embed\n",
    "        x = self.norm(self.transformer(x))\n",
    "        return self.head(x[:, 0])\n",
    "\n",
    "vit_model = SimpleViTForMedical(img_size=32).to(device)\n",
    "print(f'ViT params: {sum(p.numel() for p in vit_model.parameters()):,}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DermaScanLoss(nn.Module):\n",
    "    \"\"\"\n",
    "    TODO 5: Implement the combined loss function.\n",
    "    \n",
    "    Loss = Weighted_CE + lambda_1 * FocalLoss(melanoma) + lambda_2 * L2_reg\n",
    "    \n",
    "    Steps:\n",
    "    1. Compute weighted cross-entropy\n",
    "    2. Compute focal loss for melanoma class:\n",
    "       FL(p_t) = -(1-p_t)^gamma * log(p_t)\n",
    "    3. Return weighted combination\n",
    "    \"\"\"\n",
    "    def __init__(self, class_weights, melanoma_idx=0, focal_gamma=2.0, focal_weight=0.3):\n",
    "        super().__init__()\n",
    "        self.ce_loss = nn.CrossEntropyLoss(weight=class_weights)\n",
    "        self.melanoma_idx = melanoma_idx\n",
    "        self.focal_gamma = focal_gamma\n",
    "        self.focal_weight = focal_weight\n",
    "\n",
    "    def forward(self, logits, targets):\n",
    "        # Standard weighted cross-entropy\n",
    "        ce = self.ce_loss(logits, targets)\n",
    "\n",
    "        # Focal loss for melanoma\n",
    "        probs = F.softmax(logits, dim=-1)\n",
    "        mel_prob = probs[:, self.melanoma_idx]\n",
    "        mel_mask = (targets == self.melanoma_idx).float()\n",
    "\n",
    "        # For melanoma samples: FL = -(1-p_mel)^gamma * log(p_mel)\n",
    "        # For non-melanoma: FL = -(p_mel)^gamma * log(1-p_mel)\n",
    "        p_t = mel_prob * mel_mask + (1 - mel_prob) * (1 - mel_mask)\n",
    "        focal = -((1 - p_t) ** self.focal_gamma) * torch.log(p_t + 1e-8)\n",
    "        focal_loss = focal.mean()\n",
    "\n",
    "        return ce + self.focal_weight * focal_loss\n",
    "\n",
    "loss_fn = DermaScanLoss(class_weights.to(device))\n",
    "print('Loss function created with weighted CE + melanoma focal loss')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3.5: Training Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_evaluate(model, train_loader, val_loader, loss_fn,\n",
    "                       num_epochs=15, lr=1e-3, model_name='model'):\n",
    "    \"\"\"Train model and track melanoma-specific metrics.\"\"\"\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)\n",
    "    \n",
    "    history = {'train_loss': [], 'val_acc': [], 'mel_sensitivity': [], 'mel_auroc': []}\n",
    "    best_mel_sens = 0\n",
    "    \n",
    "    for epoch in range(num_epochs):\n",
    "        # Training\n",
    "        model.train()\n",
    "        running_loss = 0\n",
    "        for images, labels in train_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs = model(images)\n",
    "            loss = loss_fn(outputs, labels)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            running_loss += loss.item()\n",
    "        scheduler.step()\n",
    "        \n",
    "        # Validation\n",
    "        model.eval()\n",
    "        all_preds, all_labels, all_probs = [], [], []\n",
    "        with torch.no_grad():\n",
    "            for images, labels in val_loader:\n",
    "                images, labels = images.to(device), labels.to(device)\n",
    "                outputs = model(images)\n",
    "                probs = F.softmax(outputs, dim=-1)\n",
    "                _, preds = outputs.max(1)\n",
    "                all_preds.extend(preds.cpu().numpy())\n",
    "                all_labels.extend(labels.cpu().numpy())\n",
    "                all_probs.extend(probs.cpu().numpy())\n",
    "        \n",
    "        all_preds = np.array(all_preds)\n",
    "        all_labels = np.array(all_labels)\n",
    "        all_probs = np.array(all_probs)\n",
    "        \n",
    "        val_acc = (all_preds == all_labels).mean() * 100\n",
    "        \n",
    "        # Melanoma metrics\n",
    "        mel_true = (all_labels == MELANOMA_IDX).astype(int)\n",
    "        mel_pred = (all_preds == MELANOMA_IDX).astype(int)\n",
    "        mel_tp = ((mel_pred == 1) & (mel_true == 1)).sum()\n",
    "        mel_fn = ((mel_pred == 0) & (mel_true == 1)).sum()\n",
    "        mel_sensitivity = mel_tp / (mel_tp + mel_fn + 1e-8) * 100\n",
    "        \n",
    "        try:\n",
    "            mel_auroc = roc_auc_score(mel_true, all_probs[:, MELANOMA_IDX])\n",
    "        except:\n",
    "            mel_auroc = 0.5\n",
    "        \n",
    "        history['train_loss'].append(running_loss / len(train_loader))\n",
    "        history['val_acc'].append(val_acc)\n",
    "        history['mel_sensitivity'].append(mel_sensitivity)\n",
    "        history['mel_auroc'].append(mel_auroc)\n",
    "        \n",
    "        if mel_sensitivity > best_mel_sens:\n",
    "            best_mel_sens = mel_sensitivity\n",
    "            best_state = model.state_dict().copy()\n",
    "        \n",
    "        if (epoch + 1) % 5 == 0:\n",
    "            print(f'  Epoch {epoch+1}/{num_epochs}: Loss={running_loss/len(train_loader):.3f}, '\n",
    "                  f'Acc={val_acc:.1f}%, Mel Sens={mel_sensitivity:.1f}%, AUROC={mel_auroc:.3f}')\n",
    "    \n",
    "    return history, best_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train both models\n",
    "print('Training ResNet-18 Baseline...')\n",
    "baseline = create_baseline_model(img_size=32).to(device)\n",
    "loss_fn = DermaScanLoss(class_weights.to(device))\n",
    "resnet_history, resnet_best = train_and_evaluate(\n",
    "    baseline, train_loader, val_loader, loss_fn, num_epochs=15, model_name='resnet')\n",
    "\n",
    "print('\\nTraining Vision Transformer...')\n",
    "vit_model = SimpleViTForMedical(img_size=32).to(device)\n",
    "loss_fn = DermaScanLoss(class_weights.to(device))\n",
    "vit_history, vit_best = train_and_evaluate(\n",
    "    vit_model, train_loader, val_loader, loss_fn, num_epochs=15, lr=3e-4, model_name='vit')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3.6: Evaluation\n",
    "\n",
    "**TODO 6:** Compare both models on test data and compute all required metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def full_evaluation(model, test_loader, model_name='Model'):\n",
    "    \"\"\"Comprehensive evaluation with confusion matrix and ROC curves.\"\"\"\n",
    "    model.eval()\n",
    "    all_preds, all_labels, all_probs = [], [], []\n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            probs = F.softmax(model(images), dim=-1)\n",
    "            _, preds = probs.max(1)\n",
    "            all_preds.extend(preds.cpu().numpy())\n",
    "            all_labels.extend(labels.cpu().numpy())\n",
    "            all_probs.extend(probs.cpu().numpy())\n",
    "    \n",
    "    all_preds = np.array(all_preds)\n",
    "    all_labels = np.array(all_labels)\n",
    "    all_probs = np.array(all_probs)\n",
    "    \n",
    "    # Overall metrics\n",
    "    accuracy = (all_preds == all_labels).mean() * 100\n",
    "    \n",
    "    # Melanoma metrics\n",
    "    mel_true = (all_labels == MELANOMA_IDX).astype(int)\n",
    "    mel_pred = (all_preds == MELANOMA_IDX).astype(int)\n",
    "    mel_tp = ((mel_pred == 1) & (mel_true == 1)).sum()\n",
    "    mel_fp = ((mel_pred == 1) & (mel_true == 0)).sum()\n",
    "    mel_fn = ((mel_pred == 0) & (mel_true == 1)).sum()\n",
    "    mel_tn = ((mel_pred == 0) & (mel_true == 0)).sum()\n",
    "    \n",
    "    mel_sens = mel_tp / (mel_tp + mel_fn + 1e-8) * 100\n",
    "    mel_spec = mel_tn / (mel_tn + mel_fp + 1e-8) * 100\n",
    "    mel_auroc = roc_auc_score(mel_true, all_probs[:, MELANOMA_IDX])\n",
    "    \n",
    "    print(f'\\n{model_name} Results:')\n",
    "    print(f'  Overall Accuracy: {accuracy:.1f}%')\n",
    "    print(f'  Melanoma Sensitivity: {mel_sens:.1f}% (target: >90%)')\n",
    "    print(f'  Melanoma Specificity: {mel_spec:.1f}% (target: >85%)')\n",
    "    print(f'  Melanoma AUROC: {mel_auroc:.3f}')\n",
    "    \n",
    "    # Confusion matrix\n",
    "    cm = confusion_matrix(all_labels, all_preds)\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
    "    axes[0].imshow(cm, cmap='Blues')\n",
    "    axes[0].set_xticks(range(NUM_CLASSES))\n",
    "    axes[0].set_yticks(range(NUM_CLASSES))\n",
    "    axes[0].set_xticklabels(CLASS_NAMES, rotation=45, ha='right', fontsize=7)\n",
    "    axes[0].set_yticklabels(CLASS_NAMES, fontsize=7)\n",
    "    axes[0].set_title(f'{model_name} Confusion Matrix')\n",
    "    for i in range(NUM_CLASSES):\n",
    "        for j in range(NUM_CLASSES):\n",
    "            axes[0].text(j, i, str(cm[i, j]), ha='center', va='center', fontsize=8)\n",
    "    \n",
    "    # ROC curve for melanoma\n",
    "    fpr, tpr, _ = roc_curve(mel_true, all_probs[:, MELANOMA_IDX])\n",
    "    axes[1].plot(fpr, tpr, 'b-', linewidth=2, label=f'AUROC={mel_auroc:.3f}')\n",
    "    axes[1].plot([0, 1], [0, 1], 'k--', alpha=0.3)\n",
    "    axes[1].set_xlabel('False Positive Rate')\n",
    "    axes[1].set_ylabel('True Positive Rate')\n",
    "    axes[1].set_title(f'{model_name} Melanoma ROC Curve')\n",
    "    axes[1].legend()\n",
    "    axes[1].grid(alpha=0.3)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    \n",
    "    return {'accuracy': accuracy, 'mel_sensitivity': mel_sens,\n",
    "            'mel_specificity': mel_spec, 'mel_auroc': mel_auroc}\n",
    "\n",
    "# Evaluate both\n",
    "baseline.load_state_dict(resnet_best)\n",
    "resnet_metrics = full_evaluation(baseline, test_loader, 'ResNet-18')\n",
    "\n",
    "vit_model.load_state_dict(vit_best)\n",
    "vit_metrics = full_evaluation(vit_model, test_loader, 'ViT')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3.7: Error Analysis\n",
    "\n",
    "**TODO 7:** Analyze the model's failure cases, particularly false negative melanomas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def error_analysis(model, test_loader, model_name='Model'):\n",
    "    \"\"\"\n",
    "    TODO: Analyze failure cases.\n",
    "    \n",
    "    1. Find false negative melanomas (predicted non-melanoma, true melanoma)\n",
    "    2. Visualize the top 5 most confident false negatives\n",
    "    3. Compute error rate per class\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    false_negatives = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            probs = F.softmax(model(images), dim=-1)\n",
    "            preds = probs.argmax(1)\n",
    "            \n",
    "            # Find false negatives for melanoma\n",
    "            mel_mask = labels == MELANOMA_IDX\n",
    "            wrong_mask = preds != MELANOMA_IDX\n",
    "            fn_mask = mel_mask & wrong_mask\n",
    "            \n",
    "            for i in fn_mask.nonzero():\n",
    "                idx = i.item()\n",
    "                false_negatives.append({\n",
    "                    'image': images[idx].cpu(),\n",
    "                    'mel_prob': probs[idx, MELANOMA_IDX].item(),\n",
    "                    'predicted': CLASS_NAMES[preds[idx].item()]\n",
    "                })\n",
    "    \n",
    "    if false_negatives:\n",
    "        false_negatives.sort(key=lambda x: x['mel_prob'])\n",
    "        print(f'\\n{model_name}: {len(false_negatives)} false negative melanomas')\n",
    "        \n",
    "        n_show = min(5, len(false_negatives))\n",
    "        fig, axes = plt.subplots(1, n_show, figsize=(3*n_show, 3))\n",
    "        if n_show == 1: axes = [axes]\n",
    "        for i in range(n_show):\n",
    "            img = false_negatives[i]['image'].permute(1, 2, 0).numpy()\n",
    "            img = (img - img.min()) / (img.max() - img.min())\n",
    "            axes[i].imshow(img)\n",
    "            axes[i].set_title(f\"Pred: {false_negatives[i]['predicted']}\\n\"\n",
    "                             f\"Mel prob: {false_negatives[i]['mel_prob']:.3f}\", fontsize=9)\n",
    "            axes[i].axis('off')\n",
    "        plt.suptitle(f'{model_name}: Missed Melanomas', fontsize=13)\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "    else:\n",
    "        print(f'{model_name}: No false negative melanomas!')\n",
    "\n",
    "error_analysis(baseline, test_loader, 'ResNet-18')\n",
    "error_analysis(vit_model, test_loader, 'ViT')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3.8: Deployment Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def benchmark_inference(model, img_size=32, batch_size=1, num_runs=100):\n",
    "    \"\"\"Benchmark inference speed.\"\"\"\n",
    "    model.eval()\n",
    "    dummy = torch.randn(batch_size, 3, img_size, img_size).to(device)\n",
    "    \n",
    "    # Warmup\n",
    "    for _ in range(10):\n",
    "        with torch.no_grad():\n",
    "            _ = model(dummy)\n",
    "    if device.type == 'cuda':\n",
    "        torch.cuda.synchronize()\n",
    "    \n",
    "    # Benchmark\n",
    "    start = time.time()\n",
    "    for _ in range(num_runs):\n",
    "        with torch.no_grad():\n",
    "            _ = model(dummy)\n",
    "    if device.type == 'cuda':\n",
    "        torch.cuda.synchronize()\n",
    "    elapsed = time.time() - start\n",
    "    \n",
    "    latency_ms = (elapsed / num_runs) * 1000\n",
    "    throughput = num_runs / elapsed\n",
    "    print(f'  Latency: {latency_ms:.1f} ms/image')\n",
    "    print(f'  Throughput: {throughput:.0f} images/sec')\n",
    "    return latency_ms\n",
    "\n",
    "print('ResNet-18 Inference:')\n",
    "resnet_latency = benchmark_inference(baseline)\n",
    "\n",
    "print('\\nViT Inference:')\n",
    "vit_latency = benchmark_inference(vit_model)\n",
    "\n",
    "print(f'\\nViT is {vit_latency/resnet_latency:.1f}x slower than ResNet-18')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3.9: Ethical Considerations and Summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# Final comparison summary\nprint('=' * 60)\nprint('DermaScan AI: Model Comparison Summary')\nprint('=' * 60)\nprint(f'{\"Metric\":<25} {\"ResNet-18\":<15} {\"ViT\":<15} {\"Target\":<15}')\nprint('-' * 60)\nprint(f'{\"Overall Accuracy\":<25} {resnet_metrics[\"accuracy\"]:.1f}%{\"\":<10} {vit_metrics[\"accuracy\"]:.1f}%{\"\":<10} >84%')\nprint(f'{\"Mel. Sensitivity\":<25} {resnet_metrics[\"mel_sensitivity\"]:.1f}%{\"\":<10} {vit_metrics[\"mel_sensitivity\"]:.1f}%{\"\":<10} >90%')\nprint(f'{\"Mel. Specificity\":<25} {resnet_metrics[\"mel_specificity\"]:.1f}%{\"\":<10} {vit_metrics[\"mel_specificity\"]:.1f}%{\"\":<10} >85%')\nprint(f'{\"Mel. AUROC\":<25} {resnet_metrics[\"mel_auroc\"]:.3f}{\"\":<12} {vit_metrics[\"mel_auroc\"]:.3f}{\"\":<12} >0.92')\nprint(f'{\"Inference (ms)\":<25} {resnet_latency:.1f}{\"\":<12} {vit_latency:.1f}{\"\":<12} <3000')\nprint(f'{\"Parameters\":<25} {sum(p.numel() for p in baseline.parameters()):,}{\"\":<5} {sum(p.numel() for p in vit_model.parameters()):,}')\nprint('\\nDataset: DermaMNIST (real dermatoscopic images from HAM10000)')\nprint('Ethical Note: This model must be audited for performance across')\nprint('all Fitzpatrick skin types (I-VI) before clinical deployment.')\nprint('Dermoscopy AI has documented biases against darker skin tones.')"
  }
 ]
}