# -*- coding: utf-8 -*- """gradio_app_chestvision.ipynb Automatically generated by Colab. Original file is located at https://colab.research.google.com/drive/1cqkxJwlxFdpD6iRy-LNc5fhC5yoRYsw8 """ # !pip install --upgrade gradio """### Import dependencies""" import torch import torch.nn.functional as F import torchvision from torchvision import transforms, models, datasets from torch import nn, optim from torch.utils.data import DataLoader, Dataset from tqdm import tqdm from torch.utils.data import random_split import pytorch_lightning as torch_light from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint import torchmetrics from torchmetrics import Metric import os import shutil import subprocess import pandas as pd from PIL import Image import gradio from functools import partial """### Set parameters""" configs = { "IMAGE_SIZE": (224, 224), # Resize images to (W, H) "NUM_CHANNELS": 3, # RGB images "NUM_CLASSES": 15, # Number of output labels # ImageNet dataset normalization values (for pretrained backbones) "MEAN": (0.485, 0.456, 0.406), "STD": (0.229, 0.224, 0.225), "DEFAULT_BACKBONE": "ConvNeXt(tiny)", "THRESHOLD": 0.5 } BACKBONE_REGISTRY = { "ConvNeXt(small)": { "torchvision_name": "convnext_small", "ckpt": "ConvNeXt(small).ckpt"}, "ConvNeXt(tiny)": { "torchvision_name": "convnext_tiny", "ckpt": "ConvNeXt(tiny).ckpt"}, "EfficientNet(b3)": { "torchvision_name": "efficientnet_b3", "ckpt": "EfficientNet(b3).ckpt"}, "EfficientNet(v2_small)": { "torchvision_name": "efficientnet_v2_s", "ckpt": "EfficientNet(v2_small).ckpt"}, "RegNet(x3_2GF)": { "torchvision_name": "regnet_x_3_2gf", "ckpt": "RegNet(x3_2GF).ckpt"}, "ResNet50": { "torchvision_name": "resnet50", "ckpt": "ResNet50.ckpt"} } MODEL_CACHE = {} """### Define helper functions""" # helper function for loading pre-trained model # =================================================================================================== def get_pretrained_model(pretrained_model_name: str, num_classes: int, freeze_backbone: bool = True): """ Load a pretrained Torchvision classification model and replace ONLY its final Linear layer for transfer learning or fine-tuning. """ print(f"Loading pretrained [{pretrained_model_name}] model") # Load pretrained model from torchvision model = getattr(torchvision.models, pretrained_model_name)(weights="DEFAULT") # Optionally freeze all pretrained parameters (the backbone) if freeze_backbone: for param in model.parameters(): param.requires_grad = False # Get the last top-level module (typically the classifier head) from the pretrained model layer_modules = list(model.named_children()) last_module_name, last_module = layer_modules[-1] if isinstance(last_module, nn.Sequential): # If the classifier is a Sequential module, replace its final layer num_in_features = last_module[-1].in_features # find the dimensionality of the input to the last layer of the network # Replace the output layer of the classifier head to match the num. classes in the current task last_module[-1] = nn.Linear(in_features = num_in_features, out_features = num_classes) # Output layer else: # Otherwise, replace the module directly (e.g., ResNet-style fc layer) in_features = last_module.in_features setattr(model, last_module_name, nn.Linear(in_features, num_classes)) return model # helper function for preprocessing input images # =================================================================================================== preprocess_fxn = transforms.Compose( [transforms.Resize(size=configs["IMAGE_SIZE"][::-1]), transforms.ToTensor(), transforms.Normalize(configs["MEAN"], configs["STD"], inplace=True)]) # Map numeric outputs to string labels labels_dict = { 0: "Atelectasis", 1: "Cardiomegaly", 2: "Consolidation", 3: "Edema", 4: "Effusion", 5: "Emphysema", 6: "Fibrosis", 7: "Hernia", 8: "Infiltration", 9: "Mass", 10: "No finding", 11: "Nodule", 12: "Pleural_Thickening", 13: "Pneumonia", 14: "Pneumothorax"} """### Create torch lightning model (i.e., classifier) module""" class modelModule(torch_light.LightningModule): def __init__(self, num_classes = configs['NUM_CLASSES'], backbone_model_name = 'efficientnet_b3'): super().__init__() self.num_classes = num_classes self.backbone_model_name = backbone_model_name # Load a pretrained backbone and replace its final layer self.model = get_pretrained_model( num_classes = self.num_classes, pretrained_model_name = self.backbone_model_name ) # Binary classification loss operating on raw logits self.loss_function = torch.nn.BCEWithLogitsLoss() # self.accuracy_function = torchmetrics.Accuracy(task="multilabel", num_labels=self.num_classes) # self.f1_score_function = torchmetrics.F1Score(task="multilabel", num_labels=self.num_classes) self.accuracy_function = torchmetrics.classification.MultilabelAccuracy(num_labels=self.num_classes, average="weighted", threshold=0.5) self.f1_score_function = torchmetrics.classification.MultilabelF1Score(num_labels=self.num_classes, average="weighted", threshold=0.5) self.auroc_function = torchmetrics.classification.MultilabelAUROC(num_labels=self.num_classes, average="weighted", thresholds=10) self.map_score_function = torchmetrics.classification.MultilabelAveragePrecision(num_labels=self.num_classes, average="weighted", thresholds=10) # average options: macro (simple average), micro (sum), weighted (weight by class size, then avg) # threshold: Threshold for transforming probability to binary (0,1) predictions. For some metrics (e.g., AUROC), represents the number of thresholds (evenly spaced b/n 0–1) the metric should be computed at (resulting array of values are the averaged to obtain the final score) def forward(self, x): # Forward pass through the backbone model return self.model(x) def _common_step(self, batch, batch_idx): """ Shared logic for train / val / test steps. Computes loss and evaluation metrics. """ x, y = batch # Compute model predictions () y_hat = self.forward(x) # Compute metrics (expects logits + labels) # loss = self.loss_function(y_hat, y.float()) # Compute mean loss over all classes loss = torchmetrics.aggregation.MeanMetric(self.loss_function(y_hat, y.float()), weight=X.shape[0]) accuracy = self.accuracy_function(y_hat, y) f1_score = self.f1_score_function(y_hat, y) auroc = self.auroc_function(y_hat, y) mAP = self.map_score_function(y_hat, y) # mean average precision return loss, y_hat, y, accuracy, f1_score, auroc, mAP def training_step(self, batch, batch_idx): # Run shared step loss, y_hat, y, accuracy, f1_score, auroc, mAP = self._common_step(batch, batch_idx) # Log epoch-level training metrics self.log_dict( {"train_loss": loss, "train_accuracy": accuracy, "train_f1_score": f1_score, "train_auroc": auroc, "train_mAP": mAP}, on_step=False, on_epoch=True, prog_bar=True) # Lightning expects the loss key for backprop return {"loss": loss} def validation_step(self, batch, batch_idx): # Run shared step loss, y_hat, y, accuracy, f1_score, auroc, mAP = self._common_step(batch, batch_idx) # Log validation metrics self.log_dict( {"val_loss": loss, "val_accuracy": accuracy,"val_f1_score": f1_score, "val_auroc": auroc, "val_mAP": mAP}, on_step=False, on_epoch=True, prog_bar=True) def test_step(self, batch, batch_idx): # Run shared step loss, y_hat, y, accuracy, f1_score, auroc, mAP = self._common_step(batch, batch_idx) # Log test metrics self.log_dict( {"test_loss": loss, "test_accuracy": accuracy,"test_f1_score": f1_score, "test_auroc": auroc, "test_mAP": mAP}, on_step=False, on_epoch=True, prog_bar=True) def predict_step(self, batch, batch_idx): """ Prediction logic used by trainer.predict(). Returns model outputs without computing loss. """ x = batch if not isinstance(batch, (tuple, list)) else batch[0] logits = self.forward(x) # Convert logits to probabilities for inference probs = torch.sigmoid(logits) return probs def configure_optimizers(self): # Optimizer over all trainable parameters optimizer = optim.Adam(self.model.parameters(), lr=1e-3) return optimizer """### Create function for running inference (i.e., assistive medical diagnosis)""" @torch.inference_mode() def run_diagnosis( backbone_name, input_image, threshold, preprocess_fn=None, Idx2labels=None ): # Preprocess x = preprocess_fn(input_image).unsqueeze(0) # Resolve backbone backbone_info = BACKBONE_REGISTRY[backbone_name] ckpt_path = os.path.join(CKPT_ROOT, backbone_info["ckpt"]) if not os.path.exists(ckpt_path): raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}") # Load model (cache for speed) if backbone_name not in MODEL_CACHE: print(f"Loading model weights from {ckpt_path}") MODEL_CACHE[backbone_name] = modelModule.load_from_checkpoint( ckpt_path, backbone_model_name=backbone_info["torchvision_name"]) model = MODEL_CACHE[backbone_name] model.eval() # Forward logits = model(x) probs = torch.sigmoid(logits)[0].cpu().numpy() print("predicted logits\n") for i, logit_ in enumerate(logits): print(f"{Idx2labels[i]}: {logit_}") output_probs = { Idx2labels[i]: float(p) for i, p in enumerate(probs)} predicted_classes = [ Idx2labels[i] for i, p in enumerate(probs) if p >= threshold] return "\n".join(predicted_classes), output_probs """### Gradio app""" CKPT_ROOT = os.path.join(os.getcwd(), "Trained models") example_list_dir = os.path.join(os.getcwd(), "Curated test samples") example_list_img_names = os.listdir(example_list_dir) # example_list = [ # [os.path.join(example_list_dir, example_img), configs["DEFAULT_BACKBONE"]] # for example_img in example_list_img_names # if example_img.lower().endswith(".png")] example_list = [ [configs["DEFAULT_BACKBONE"], os.path.join(example_list_dir, example_img)] for example_img in example_list_img_names[:8] if example_img.lower().endswith(".png")] # example_list = [['/content/new_labels.csv',"ResNet50"]] gradio_app = gradio.Interface( fn = partial(run_diagnosis, preprocess_fn = preprocess_fxn, Idx2labels = labels_dict), inputs = [gradio.Dropdown(list(BACKBONE_REGISTRY.keys()), value="ResNet50", label="Select Backbone Model"), gradio.Image(type="pil", label="Load chest-X-ray image here"), gradio.Slider(minimum = 0.1, maximum = 0.9, step = 0.05, value = 0.4, label = "Set Prediction Threshold") ], outputs = [gradio.Textbox(label="Predicted Medical Condition(s)"), gradio.Label(label="Predicted Probabilities", show_label=False)], examples = example_list, cache_examples = False, title = "ChestVision", description = "Deep CNN-based solutions for assistive medical diagnosis", article = "Author: C. Foli (02.2026) | Website: coming soon...") if __name__ == "__main__": gradio_app.launch()