Update app.py

app.py

@@ -1,15 +1,916 @@
+import os, re, random, json
 import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+
+from torch.optim import Optimizer, Adam, AdamW, SGD
+from transformers import AutoTokenizer, AutoModel, get_linear_schedule_with_warmup
+from datasets import load_dataset
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics import classification_report, confusion_matrix, f1_score
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+import seaborn as sns
+import timm
+import torchvision
+import torchvision.transforms as transforms
+from torchvision.datasets import  ImageFolder
+from collections import Counter
+from torchvision.transforms import ToTensor
+from torchmetrics import MeanMetric, Accuracy
+from torchmetrics import ConfusionMatrix, Accuracy, Precision, Recall, F1Score
+
+from sklearn.model_selection import train_test_split
+from pathlib import Path
+
+from tqdm import tqdm
+from typing import Tuple, Union, Callable
+
 import gradio as gr
 
-def sepia(input_img):
-    sepia_filter = np.array([
-        [0.393, 0.769, 0.189],
-        [0.349, 0.686, 0.168],
-        [0.272, 0.534, 0.131]
-    ])
-    sepia_img = input_img.dot(sepia_filter.T)
-    sepia_img /= sepia_img.max()
-    return sepia_img
-
-demo = gr.Interface(sepia, gr.Image(), "image", api_name="predict")
-demo.launch()
+from PIL import Image
+# Ensure dataset outputs images to the same size
+# because the model expects the input to be consistent
+from torchvision.transforms import v2
+import random
+from glob import glob
+
+
+class CNN32(nn.Module):
+    """
+    32-layer CNN - 4 blocks of 8 conv layers each.
+    Channels:  Block1 = 32  Block2 = 64  Block3 = 128  Block4 = 256
+    Spatial:   224 -> 112 -> 56 -> 28 -> 14 -> AdaptiveAvgPool(1)
+
+    BatchNorm added after every conv essential bcause
+    gradients vanish through 32 layers and the network
+    fails to train (this is exactly why ResNet was invented).
+
+    https://d2l.ai/chapter_convolutional-modern/batch-norm.html
+    https://docs.pytorch.org/docs/2.12/generated/torch.nn.BatchNorm2d.html
+    https://arxiv.org/pdf/1502.03167
+    """
+
+    @property
+    def has_backbone(): return False
+
+    def __init__(self, num_classes=8):
+        super().__init__()
+
+        self.features = nn.Sequential(
+            # Block 1: 8 convolution & pooling layers, 32 channels
+            nn.LazyConv2d(32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.LazyConv2d(32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.LazyConv2d(32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.LazyConv2d(32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.LazyConv2d(32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.LazyConv2d(32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.LazyConv2d(32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.LazyConv2d(32, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.MaxPool2d(kernel_size=2),          # from 224 to 112
+            nn.Dropout2d(0.1),
+
+            # Block 1: 8 convolution & pooling layers, 64 channels
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.MaxPool2d(kernel_size=2),          # from 112 to 56
+            nn.Dropout2d(0.2),
+
+            # Block 3: 8 convolution & pooling layers, 128 channels
+            nn.LazyConv2d(128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+
+            nn.LazyConv2d(128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+
+            nn.LazyConv2d(128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+
+            nn.LazyConv2d(128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+
+            nn.LazyConv2d(128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+
+            nn.LazyConv2d(128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+
+            nn.LazyConv2d(128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+
+            nn.LazyConv2d(128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+
+            nn.MaxPool2d(kernel_size=2),          # from 56 to 28
+            nn.Dropout2d(0.2),
+
+            # Block 4: 8 convolution & pooling layers, 256 channels
+            nn.LazyConv2d(256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+
+            nn.LazyConv2d(256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+
+            nn.LazyConv2d(256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+
+            nn.LazyConv2d(256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+
+            nn.LazyConv2d(256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+
+            nn.LazyConv2d(256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+
+            nn.LazyConv2d(256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+
+            nn.LazyConv2d(256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+
+            nn.MaxPool2d(kernel_size=2),          # from 28 to 14
+            nn.Dropout2d(0.3),
+        )
+
+        # Collapses 14 x 14 -> 1 x 1 regardless of input size
+        # https://discuss.pytorch.org/t/what-is-adaptiveavgpool2d/26897
+        # https://medium.com/@caring_smitten_gerbil_914/demystifying-nn-adaptiveavgpool2d-in-pytorch-why-adaptive-pooling-matters-in-deep-learning-1f7b7b1cc9b0
+        # https://docs.pytorch.org/docs/main/generated/torch.nn.modules.pooling.AdaptiveAvgPool2d.html
+        self.pool = nn.AdaptiveAvgPool2d(1)
+
+        # Fully Connected Layer head: wider than CNN9 to match 256 input channels
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.LazyLinear(1024),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.LazyLinear(512),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.LazyLinear(256),
+            nn.ReLU(),
+            nn.LazyLinear(num_classes),  # no Dropout on output
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.pool(x)
+        return self.classifier(x)
+
+
+
+
+class ResNet50Modified(nn.Module):
+  """
+  https://pytorch.org/hub/nvidia_deeplearningexamples_resnet50/
+  https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html
+  https://medium.com/@deepvisionkararhaider/resnet-50-explained-step-by-step-the-easiest-guide-to-deep-residual-networks-7616f4f45046
+  https://arxiv.org/pdf/1512.03385
+  """
+  @property
+  def has_backbone(): return True
+
+  def __init__(self, num_classes=8):
+      super().__init__()
+      self.backbone  = torchvision.models.resnet50(weights="IMAGENET1K_V2", progress=True)
+      in_features  = self.backbone.fc.in_features   # 2048 the standard feature embedding vector size
+
+      # replace the original fc layer (a Linear(2048, 1000) trained on ImageNet's 1000 classes) with a passthrough.
+      # The backbone now outputs the raw 2048-dimensional feature vector instead of 1000 class logits.
+      self.backbone.fc = nn.Identity()
+
+      # Fully conncted layer head outside self.backbone -> stays trainable during phase 1 freeze
+      self.classifier = nn.Sequential(
+          nn.Dropout(p=0.4),
+          nn.Linear(in_features, num_classes),
+      )
+
+  def forward(self, x):
+      return self.classifier(self.backbone(x))
+
+
+
+
+class ResNet18Modified(nn.Module):
+  """
+  https://pytorch.org/hub/nvidia_deeplearningexamples_resnet50/
+  https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html
+  https://medium.com/@deepvisionkararhaider/resnet-50-explained-step-by-step-the-easiest-guide-to-deep-residual-networks-7616f4f45046
+  https://arxiv.org/pdf/1512.03385
+  """
+
+  @property
+  def has_backbone(): return True
+
+  def __init__(self, num_classes=8):
+      super().__init__()
+      self.backbone  = torchvision.models.resnet18(weights="IMAGENET1K_V1", progress=True)
+      in_features  = self.backbone.fc.in_features   # 2048 the standard feature embedding vector size
+
+      # replace the original fc layer (a Linear(2048, 1000) trained on ImageNet's 1000 classes) with a passthrough.
+      # The backbone now outputs the raw 2048-dimensional feature vector instead of 1000 class logits.
+      self.backbone.fc = nn.Identity()
+
+      # Fully conncted layer head outside self.backbone -> stays trainable during phase 1 freeze
+      self.classifier = nn.Sequential(
+          nn.Dropout(p=0.4),
+          nn.Linear(in_features, num_classes),
+      )
+
+  def forward(self, x):
+      return self.classifier(self.backbone(x))
+
+
+
+class DenseNet201Modified(nn.Module):
+    """
+     Densenet201: https://docs.pytorch.org/vision/main/models/generated/torchvision.models.densenet201.html#torchvision.models.densenet201
+     https://medium.com/@karuneshu21/implement-densenet-in-pytorch-46374ef91900
+     https://docs.pytorch.org/vision/main/models/densenet.html
+     Densely Connected Convolutional Networks: https://arxiv.org/abs/1608.06993
+    """
+
+    @property
+    def has_backbone(): return True
+
+    def __init__(self, num_classes=8):
+        super().__init__()
+        self.backbone = torchvision.models.densenet201(
+            weights="IMAGENET1K_V1", progress=True)
+        in_features = self.backbone.classifier.in_features
+        self.backbone.classifier = nn.Identity()
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=0.4),
+            nn.Linear(in_features, num_classes),
+        )
+    def forward(self, x):
+        return self.classifier(self.backbone(x))
+
+
+
+
+class XceptionModified(nn.Module):
+    """
+
+    https://huggingface.co/docs/timm/en/models/xception
+    You can finetune any of the pre-trained models just by changing the classifier (the last layer).
+    """
+
+    @property
+    def has_backbone(): return True
+
+    def __init__(self, num_classes=8):
+        super().__init__()
+        self.backbone = timm.create_model(
+            "xception", pretrained=True, num_classes=0)
+        in_features   = self.backbone.num_features
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=0.4),
+            nn.Linear(in_features, num_classes),
+        )
+    def forward(self, x):
+        return self.classifier(self.backbone(x))
+
+
+
+class InceptionV3Modified(nn.Module):
+  """
+  Special case:  requires 299×299 input
+  https://arxiv.org/abs/1512.00567
+  https://docs.pytorch.org/vision/main/models/generated/torchvision.models.inception_v3.html
+  """
+
+  @property
+  def has_backbone(): return True
+
+  def __init__(self, num_classes=8):
+      super().__init__()
+      self.backbone = torchvision.models.inception_v3(
+          weights="IMAGENET1K_V1", progress=True,
+          aux_logits=True)
+      self.backbone.aux_logits = False  # disable after loading -> forward returns plain tensor
+      self.backbone.AuxLogits = None         # free the auxiliary classifier modules
+
+      in_features = self.backbone.fc.in_features  # 2048 channels
+      self.backbone.fc = nn.Identity()
+      self.classifier  = nn.Sequential(
+          nn.Dropout(p=0.4),
+          nn.Linear(in_features, num_classes),
+      )
+  def forward(self, x):
+      return self.classifier(self.backbone(x))
+
+
+
+
+class AlexNetModified(nn.Module):
+    """
+    AlexNet (Krizhevsky et al. 2012) - historical baseline:
+    https://www.researchgate.net/publication/319770183_Imagenet_classification_with_deep_convolutional_neural_networks.
+    https://proceedings.neurips.cc/paper_files/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf
+    First deep CNN to win ImageNet. Shows progression from early architectures.
+    weights="IMAGENET1K_V1"
+    https://medium.com/@shivsingh483/understanding-alexnet-the-2012-breakthrough-that-changed-ai-forever-7c365cf76969
+    https://docs.pytorch.org/vision/main/models/generated/torchvision.models.alexnet.html
+    """
+    @property
+    def has_backbone(): return True
+
+    def __init__(self, num_classes=8):
+        super().__init__()
+        backbone = torchvision.models.alexnet(weights="IMAGENET1K_V1")
+
+        # Remove the final linear classifier
+        self.backbone = nn.Sequential(backbone.features, backbone.avgpool, nn.Flatten(),
+                                  *list(backbone.classifier.children())[:-1]) # igore the last layer
+        in_features = 4096
+
+        self.classifier = nn.Sequential(
+          nn.Dropout(p=0.4),
+          nn.Linear(in_features, num_classes),
+      )
+
+    def forward(self, x):
+        return self.classifier(self.backbone(x))
+
+
+
+
+class GoogLeNetModified(nn.Module):
+    """
+    GoogLeNet
+
+    use weights="IMAGENET1K_V1" and aux_logits=False for simplicity.
+    https://www.researchgate.net/publication/316215961_KVASIR_A_Multi-Class_Image_Dataset_for_Computer_Aided_Gastrointestinal_Disease_Detection
+    https://doras.dcu.ie/21821/1/Pogorelov_et_al._2017.pdf
+    https://arxiv.org/pdf/1409.4842
+    https://pytorch.org/hub/pytorch_vision_googlenet/
+    https://pytorch.org/hub/pytorch_vision_googlenet/
+    https://www.google.com/url?sa=t&source=web&rct=j&opi=89978449&url=https://medium.com/%40siddheshb008/googlenet-a-deep-dive-into-googles-neural-network-technology-f588d1b49e55&ved=2ahUKEwjdirymi76UAxV0VEEAHR-aGQQQFnoECCUQAQ&usg=AOvVaw1Bij1bxrw7vGia5oJQhiZh
+    https://www.cs.unc.edu/~wliu/papers/GoogLeNet.pdf
+
+     # add final nn.Linear classifier layer
+    """
+
+    @property
+    def has_backbone(): return True
+
+    def __init__(self, num_classes=8):
+        super().__init__()
+        self.backbone = torchvision.models.googlenet(
+            weights="IMAGENET1K_V1",
+            aux_logits=True,          # required by torchvision when loading weights
+        )
+        self.backbone.aux_logits = False  # disable after loading -> forward returns plain tensor
+        self.backbone.aux1 = None         # free the auxiliary classifier modules
+        self.backbone.aux2 = None
+        self.backbone.AuxLogits = None
+
+        in_features = self.backbone.fc.in_features
+        self.backbone.fc = nn.Identity()  # strip head from backbone
+
+        self.classifier = nn.Sequential(
+            nn.Dropout(0.4),
+            nn.Linear(in_features, num_classes),
+        )
+
+    def forward(self, x):
+        return self.classifier(self.backbone(x))
+
+
+
+class VGG16Modified(nn.Module):
+    """
+    VGG16 (Simonyan & Zisserman 2014).
+    https://arxiv.org/abs/1409.1556
+    https://arxiv.org/pdf/1409.1556
+    https://www.robots.ox.ac.uk/~vgg/research/very_deep/
+
+    You can finetune any of the pre-trained models just by changing the classifier (the last layer).
+    # num_classes=0 -> remove classifier nn.Linear
+    """
+
+    @property
+    def has_backbone(): return True
+
+    def __init__(self, num_classes=8):
+        super().__init__()
+        self.backbone = torchvision.models.vgg16(weights="IMAGENET1K_V1")
+
+        # Remove last linear classifier
+        in_features = self.backbone.classifier[6].in_features  # 4096
+        self.backbone.classifier[6] = nn.Identity() # strip original classifier head
+
+
+        self.classifier = nn.Sequential(
+            nn.Dropout(0.4),
+            nn.Linear(in_features, num_classes),
+        )
+
+    def forward(self, x):
+        return self.classifier(self.backbone(x))
+
+
+
+
+class EfficientNetB0Modified(nn.Module):
+    """
+    EfficientNet-B0 (Tan & Le 2019)
+
+    https://proceedings.mlr.press/v97/tan19a.html
+    https://arxiv.org/pdf/1905.11946
+     # num_classes=0 -> remove classifier nn.Linear
+    """
+
+    @property
+    def has_backbone(): return True
+
+    def __init__(self, num_classes=8):
+        super().__init__()
+        self.backbone = timm.create_model("efficientnet_b0", pretrained=True, num_classes=0)
+        self.classifier = nn.Sequential(
+            nn.Dropout(0.4),
+            nn.Linear(self.backbone.num_features, num_classes),
+        )
+
+    def forward(self, x):
+        return self.classifier(self.backbone(x))
+
+
+class SEBlock(nn.Module):
+    """
+    Squeeze-and-Excitation block.
+
+    https://www.emergentmind.com/topics/squeeze-and-excitation-se-mechanism
+    https://www.digitalocean.com/community/tutorials/channel-attention-squeeze-and-excitation-networks
+    https://arxiv.org/pdf/1709.01507
+
+    Learns WHICH feature channels matter most for each endoscopic finding.
+    For Example: colour channels matter more for esophagitis, texture for polyps.
+    """
+    def __init__(self, channels, reduction=16):
+        super().__init__()
+        self.pool   = nn.AdaptiveAvgPool2d(1)
+        self.excite = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(channels, channels // reduction, bias=False),
+            nn.ReLU(),
+            nn.Linear(channels // reduction, channels, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        s = self.excite(self.pool(x)).unsqueeze(-1).unsqueeze(-1)
+        return x * s
+
+class EfficientNetB0_SE_Modified(nn.Module):
+    """
+    EfficientNet-B0 with custom SE attention pooling.
+
+    Modification over baseline EfficientNet-B0:
+    - After the backbone's final feature maps, apply an SE block
+      that recalibrates channel importance before classification.
+
+    https://arxiv.org/pdf/1905.11946
+    https://medium.com/codex/a-summary-of-efficientnet-rethinking-model-scaling-for-cnns-d524d37ff8bb
+
+     # num_classes=0 -> remove classifier nn.Linear
+    """
+
+    @property
+    def has_backbone(): return True
+
+    def __init__(self, num_classes=8, reduction=16):
+        super().__init__()
+        self.backbone = timm.create_model(
+            "efficientnet_b0",
+            pretrained=True,
+            num_classes=0,
+            global_pool=""
+        )
+        in_features = self.backbone.num_features   # 1280
+        self.se_block   = SEBlock(in_features, reduction=reduction)
+        self.pool       = nn.AdaptiveAvgPool2d(1)      # [B, 1280, H, W] -> [B, 1280, 1, 1]
+
+        self.classifier = nn.Sequential(
+            nn.Flatten(),              #  [B, 1280, 1, 1] -> [B, 1280]
+            nn.Dropout(0.4),
+            nn.Linear(in_features, num_classes),
+        )
+
+    def forward(self, x):
+        x = self.backbone(x)           # [B, 1280, H, W]
+        x = self.se_block(x)           # [B, 1280, H, W]
+        x = self.pool(x)               # [B, 1280, 1, 1]
+        return self.classifier(x)      # Flatten inside classifier -> [B, 8]
+
+
+
+
+# must have __init__ and forward
+class MaxVitTinyTfModified(nn.Module):
+  """
+  https://huggingface.co/timm/maxvit_tiny_tf_224.in1k
+  You can finetune any of the pre-trained models just by changing the classifier (the last layer).
+  """
+
+  @property
+  def has_backbone(): return True
+
+  def __init__(self, num_classes=8): # define different parts of the model
+      super().__init__()
+      self.backbone = timm.create_model(
+          "maxvit_tiny_tf_224", pretrained=True,
+          num_classes=0  # remove classifier nn.Linear
+      )
+      out_size = self.backbone.num_features
+      self.classifier = nn.Sequential(
+          nn.Dropout(p=0.4),
+          nn.Linear(out_size, num_classes)
+      )
+
+   # take example or batch of examples and connect the parts
+   # defind in init and return the output
+  def forward(self, x):
+      x = self.backbone(x)   # returns pooled features, no head
+      return self.classifier(x)
+
+
+
+
+from dataclasses import dataclass, field
+import torch
+
+@dataclass(frozen=True)
+class EvaluationResult:
+  confusion_matrix: ConfusionMatrix
+  accuracy: float
+  precision: float
+  recall: float
+  f1_score: float
+  all_preds: list = field(default_factory=list)
+  all_labels: list = field(default_factory=list)
+
+
+
+# Model List
+MODEL_REGISTRY_ADAM = {
+    #"SimpleCNN_Adam": (SimpleCNN(NUM_CLASSES), False),  # no base model
+    #"CNN3_Adam": (CNN3(NUM_CLASSES), False),  # no base model
+    #"CNN6_Adam": (CNN6(NUM_CLASSES), False),  # no base model
+    #"CNN9_Adam": (CNN9(NUM_CLASSES), False),  # no base model
+    'CNN32_Adam': (CNN32(NUM_CLASSES), False),
+    "Resnet50_Adam": (ResNet50Modified(NUM_CLASSES), True), # My Method 1
+    "Resnet18_Adam": (ResNet18Modified(NUM_CLASSES), True),
+    #"AlexNet_Adam": (AlexNetModified(NUM_CLASSES), True),
+    #"GoogLeNet_Adam":(GoogLeNetModified(NUM_CLASSES), True),
+    #"VGG16_Adam": (VGG16Modified(NUM_CLASSES), True),
+    #"EfficientNet-B0_Adam": (EfficientNetB0Modified(NUM_CLASSES), True),
+    "EfficientNet-B0-SE_Adam":(EfficientNetB0_SE_Modified(NUM_CLASSES), True),  # My Method 2
+    #"MaxVitTinyTf_Adam":(MaxVitTinyTfModified(NUM_CLASSES), True),
+    #'Xception_Adam': (XceptionModified(NUM_CLASSES), True),
+}
+
+special_models_adam = {
+    'InceptionV3_Adam': (InceptionV3Modified(NUM_CLASSES), True)
+}
+
+# Model List
+MODEL_REGISTRY_SGD = {
+    #"SimpleCNN_SGD": (SimpleCNN(NUM_CLASSES), False),  # no base model
+    #"CNN3_SGD": (CNN3(NUM_CLASSES), False),  # no base model
+    #"CNN6_SGD": (CNN6(NUM_CLASSES), False),  # no base model
+    #"CNN9_SGD": (CNN9(NUM_CLASSES), False),  # no base model
+    'CNN32_SGD': (CNN32(NUM_CLASSES), False),
+    "Resnet50_SGD": (ResNet50Modified(NUM_CLASSES), True), # My Method 1
+    "Resnet18_SGD": (ResNet18Modified(NUM_CLASSES), True),
+    #"AlexNet_SGD": (AlexNetModified(NUM_CLASSES), True),
+    #"GoogLeNet_SGD":(GoogLeNetModified(NUM_CLASSES), True),
+    #"VGG16_SGD": (VGG16Modified(NUM_CLASSES), True),
+    #"EfficientNet-B0_SGD": (EfficientNetB0Modified(NUM_CLASSES), True),
+    "EfficientNet-B0-SE_SGD":(EfficientNetB0_SE_Modified(NUM_CLASSES), True),  # My Method 2
+    #"MaxVitTinyTf_SGD":(MaxVitTinyTfModified(NUM_CLASSES), True),
+    #'Xception_SGD': (XceptionModified(NUM_CLASSES), True),
+}
+
+special_models_sgd = {
+    'InceptionV3_SGD': (InceptionV3Modified(NUM_CLASSES), True)
+}
+
+
+
+MODEL_REGISTRY_ADAM_W = {
+    #"SimpleCNN_AdamW": (SimpleCNN(NUM_CLASSES), False),  # no base model
+    #"CNN3_AdamW": (CNN3(NUM_CLASSES), False),  # no base model
+    #"CNN6_AdamW": (CNN6(NUM_CLASSES), False),  # no base model
+    #"CNN9_AdamW": (CNN9(NUM_CLASSES), False),  # no base model
+    'CNN32_AdamW': (CNN32(NUM_CLASSES), False),
+    "Resnet50_AdamW": (ResNet50Modified(NUM_CLASSES), True), # My Method 1
+    "Resnet18_AdamW": (ResNet18Modified(NUM_CLASSES), True),
+    #"AlexNet_AdamW": (AlexNetModified(NUM_CLASSES), True),
+    #"GoogLeNet_AdamW":(GoogLeNetModified(NUM_CLASSES), True),
+    #"VGG16_AdamW": (VGG16Modified(NUM_CLASSES), True),
+    #"EfficientNet-B0_AdamW": (EfficientNetB0Modified(NUM_CLASSES), True),
+    "EfficientNet-B0-SE_AdamW":(EfficientNetB0_SE_Modified(NUM_CLASSES), True),  # My Method 2
+    #"MaxVitTinyTf_AdamW":(MaxVitTinyTfModified(NUM_CLASSES), True),
+    #'Xception_AdamW': (XceptionModified(NUM_CLASSES), True),
+}
+
+special_models_adam_w = {
+    'InceptionV3_AdamW': (InceptionV3Modified(NUM_CLASSES), True)
+}
+
+
+best_optimizer = 'AdamW'
+
+
+
+LOADED_MODELS = {}
+
+if best_optimizer == "SGD":
+  for model_name, (model, _) in {**MODEL_REGISTRY_SGD, **special_models_sgd}.items():
+    ckpt_path = f"best_{model_name.replace(' ','_')}.pth"
+    try:
+        model.load_state_dict(
+            torch.load(ckpt_path, map_location=DEVICE))
+        model.to(DEVICE)
+        model.eval()
+        LOADED_MODELS[model_name] = model
+        print(f" ✔ Loaded: {model_name}")
+    except FileNotFoundError:
+        print(f"(<-->) Skipped: {model_name} — checkpoint not found ({ckpt_path})")
+elif best_optimizer == "AdamW":
+  for model_name, (model, _) in {**MODEL_REGISTRY_ADAM_W, **special_models_adam_w}.items():
+    ckpt_path = f"best_{model_name.replace(' ','_')}.pth"
+    try:
+        model.load_state_dict(
+            torch.load(ckpt_path, map_location=DEVICE))
+        model.to(DEVICE)
+        model.eval()
+        LOADED_MODELS[model_name] = model
+        print(f" ✔ Loaded: {model_name}")
+    except FileNotFoundError:
+        print(f"(<-->) Skipped: {model_name} — checkpoint not found ({ckpt_path})")
+elif best_optimizer == "Adam":
+  for model_name, (model, _) in {**MODEL_REGISTRY_ADAM, **special_models_adam}.items():
+    ckpt_path = f"best_{model_name.replace(' ','_')}.pth"
+    try:
+        model.load_state_dict(
+            torch.load(ckpt_path, map_location=DEVICE))
+        model.to(DEVICE)
+        model.eval()
+        LOADED_MODELS[model_name] = model
+        print(f" ✔ Loaded: {model_name}")
+    except FileNotFoundError:
+        print(f"(<-->) Skipped: {model_name} — checkpoint not found ({ckpt_path})")
+
+print(f"\nAvailable models: {list(LOADED_MODELS.keys())}")
+
+
+inference_transform = v2.Compose([
+    v2.Resize((224, 224)),
+    v2.ToImage(),
+    v2.ToDtype(torch.float32, scale=True),
+    v2.Normalize(mean=[0.485, 0.456, 0.406],
+                 std=[0.229, 0.224, 0.225]),
+])
+
+
+inception_transform = v2.Compose([
+    v2.Resize((299, 299)),
+    v2.ToImage(),
+    v2.ToDtype(torch.float32, scale=True),
+    v2.Normalize(mean=[0.485, 0.456, 0.406],
+                 std=[0.229, 0.224, 0.225]),
+])
+
+CLASS_NAMES = [
+    "dyed-lifted-polyps",
+    "dyed-resection-margins",
+    "esophagitis",
+    "normal-cecum",
+    "normal-pylorus",
+    "normal-z-line",
+    "polyps",
+    "ulcerative-colitis",
+]
+
+
+CLASS_EMOJIS = {
+    "dyed-lifted-polyps":      "🟣",
+    "dyed-resection-margins":  "🔵",
+    "esophagitis":             "🟠",
+    "normal-cecum":            "🟢",
+    "normal-pylorus":          "🟨",
+    "normal-z-line":           "🟦",
+    "polyps":                  "🟡",
+    "ulcerative-colitis":      "🔴",
+}
+
+
+# https://www.gradio.app/guides/the-interface-class
+# fn: the function to wrap a user interface (UI) around
+def predict(image, model_name: str):
+    """
+    Takes a PIL image and model name.
+    Returns: prediction label, confidence bar chart figure, attention note.
+    """
+     # Guard: Image is None
+    if image is None:
+        return "No image uploaded.", None
+
+    # Guad: if model_name is None
+    if model_name is None:
+        return "No model selected.", None
+
+    # Guard: Model not loaded
+    if model_name not in LOADED_MODELS:
+        return f"Model '{model_name}' not loaded. Available: {list(LOADED_MODELS.keys())}", None
+
+    model = LOADED_MODELS[model_name]
+
+    # Guard: In case model is stored as None
+    if model is None:
+        return f"Model '{model_name}' is None — checkpoint failed to load.", None
+
+    # Image Transformation
+    transform = (inception_transform
+                 if "inception" in model_name.lower()
+                 else inference_transform)
+
+    img_tensor = transform(image).unsqueeze(0).to(DEVICE)  # [1, 3, H, W]
+
+    # Model Prediction Inference
+    model.eval()
+    with torch.no_grad():
+        logits     = model(img_tensor)
+        probs      = F.softmax(logits, dim=1).squeeze()
+        pred_idx   = probs.argmax().item()
+        pred_class = CLASS_NAMES[pred_idx]
+        confidence = probs[pred_idx].item()
+
+    # Probability Bar Chart
+    probs_np = probs.cpu().numpy()
+    colours  = ["#1D9E75" if i == pred_idx else "#B0BEC5"
+                for i in range(len(CLASS_NAMES))]
+
+    fig, ax = plt.subplots(figsize=(7, 3.5))
+    bars = ax.barh(CLASS_NAMES, probs_np, color=colours, edgecolor="white")
+    ax.set_xlim(0, 1)
+    ax.set_xlabel("Probability")
+    ax.set_title(f"{model_name} — class probabilities")
+    ax.spines[["top", "right"]].set_visible(False)
+    for bar, prob in zip(bars, probs_np):
+        if prob > 0.02:
+            ax.text(prob + 0.01, bar.get_y() + bar.get_height() / 2,
+                    f"{prob:.1%}", va="center", fontsize=9)
+    plt.tight_layout()
+
+    # Class Label
+    emoji = CLASS_EMOJIS.get(pred_class, "")
+    label = (f"{emoji} Predicted: {pred_class.replace('-',' ').title()}\n"
+             f"Confidence: {confidence:.1%}\n"
+             f"Model: {model_name}")
+
+    return label, fig
+
+
+print("LOADED_MODELS contents:")
+for name, model in LOADED_MODELS.items():
+    print(f"  {name}: {type(model).__name__ if model is not None else 'None'}")
+
+example_paths = []
+for class_name in CLASS_NAMES:
+    images = glob(f"gradio_examples/{class_name}.jpg")
+    if images:
+        dst = f"gradio_examples/{class_name}.jpg"
+        example_paths.append(dst)
+
+
+with gr.Blocks(title="Kvasir GI Endoscopy Multi-Image Classifier") as demo:
+
+    gr.Markdown("""
+    # Kvasir Gastrointestinal Endoscopy Classifier
+    Upload an endoscopy image and select a model to classify it into one of
+    **8 GI tract categories**: esophagitis, polyps, ulcerative colitis,
+    dyed lifted polyps, dyed resection margins, normal cecum, pylorus, or z-line.
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(
+                type="pil", # PIL Image Format
+                label="Upload endoscopy image",
+            )
+            model_dropdown = gr.Dropdown(
+                choices=list(LOADED_MODELS.keys()),
+                value=list(LOADED_MODELS.keys())[0],
+                label="Select model",
+            )
+            predict_btn = gr.Button("Classify", variant="primary")
+
+        with gr.Column(scale=2):
+            label_output = gr.Textbox(
+                label="Prediction",
+                lines=3,
+            )
+            chart_output = gr.Plot(
+                label="Class probabilities",
+            )
+
+    # Example images for testing
+    gr.Examples(
+        examples=[[p, random.choice(list(LOADED_MODELS.keys()))] for p in example_paths],
+        inputs=[image_input, model_dropdown],
+        label="Example images",
+    )
+
+    # Wire button to predict function
+    predict_btn.click(
+        fn=predict,
+        inputs=[image_input, model_dropdown],
+        outputs=[label_output, chart_output],
+    )
+
+    # Also predict on image upload (no button press needed)
+    image_input.change(
+        fn=predict,
+        inputs=[image_input, model_dropdown],
+        outputs=[label_output, chart_output],
+    )
+
+demo.launch(share=True, allowed_paths=["gradio_examples"])   # share=True -> public URL