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