app.py

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

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


DEVICE = "cuda" if torch.cuda.is_available() \
          else "mps" if torch.mps.is_available() \
          else "cpu"

print("Device: ", DEVICE)


CLASS_NAMES = [
    "dyed-lifted-polyps",
    "dyed-resection-margins",
    "esophagitis",
    "normal-cecum",
    "normal-pylorus",
    "normal-z-line",
    "polyps",
    "ulcerative-colitis",
]

NUM_CLASSES = len(CLASS_NAMES)


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"models/15_epochs/best_{model_name.replace(' ','_')}.pth"
    try:
        model.load_state_dict(
            torch.load(ckpt_path, map_location=DEVICE, weights_only=False))
        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"models/15_epochs/best_{model_name.replace(' ','_')}.pth"
    try:
        model.load_state_dict(
            torch.load(ckpt_path, map_location=DEVICE, weights_only=False))
        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"models/15_epochs/best_{model_name.replace(' ','_')}.pth"
    try:
        model.load_state_dict(
            torch.load(ckpt_path, map_location=DEVICE, weights_only=False))
        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_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