app.py

# -*- coding: utf-8 -*-
"""gradio_app_chestvision-PRO.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1gVrx5TyipNPvn8D7GaK0pNBCnLeYTAD_
"""

"""### 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
from transformers import AutoModel, pipeline

"""### Initialize Containers"""

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": "CheXFormer-small",
    "DEFAULT_VLM": "Lingshu-7B",

    "THRESHOLD": 0.2
    }

ViT_REGISTRY = {
    "CheXFormer-small": "m42-health/CXformer-small",
    # "CheXFormer-base": "m42-health/CXformer-base",
    "ViT-base-16": "google/vit-base-patch16-224"}

VLM_REGISTRY = {
    "MedMO": "MBZUAI/MedMO-8B",
    "Qwen3-VL-2B": "Qwen/Qwen3-VL-2B-Instruct",
    "Lingshu-7B": "lingshu-medical-mllm/Lingshu-7B",
    "MedGemma-4b": "google/medgemma-1.5-4b-it"}

VLM_SYSTEM_PROMPT = """ You are a medical imaging assistant specializing in chest radiography.

      A trained multi-label classifier analyzed a chest X-ray and made a prediction, including predicted medical condition(s) and their associated probabilities:

      Your task:
      1. Analyze the chest X-ray image to identify key features supporting the predicted condition(s).
      2. Do NOT introduce new diagnoses.
      3. Only explain radiographic findings that could support the listed prediction(s).
      4. Use cautious, uncertainty-aware language.
      5. If probability < 0.50, emphasize uncertainty.
      6. Do NOT contradict the classifier.

      Structure your answer as:

      Observed Radiographic Findings:
      ...

      How Chest X-ray Features Support the Predicted Conditions:
      ...
      """

ViT_MODEL_CACHE = {}
VLM_MODEL_CACHE = {}

"""### Define helper functions"""

# helper function for loading pre-trained model
# ===================================================================================================
class get_pretrained_model(nn.Module):
    def __init__(
        self,
        model_name: str,
        num_classes: int,
        num_layers_to_unfreeze: int = 0):
        super().__init__()

        print(f"Loading pretrained [{model_name}] model")

        self.backbone = AutoModel.from_pretrained(
            ViT_REGISTRY[model_name],
            # model_name,
            trust_remote_code=True)

        hidden_size = self.backbone.config.hidden_size

        # Freeze entire backbone first
        for param in self.backbone.parameters():
            param.requires_grad = False

        # Selectively unfreeze last N layers
        if num_layers_to_unfreeze > 0:
            self._unfreeze_last_n_layers(num_layers_to_unfreeze)

        # Single classification head
        self.classifier = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Dropout(0.4),
            nn.Linear(hidden_size, num_classes) )

    def forward(self, x):
        outputs = self.backbone(x)

        # Use CLS token
        img_embeddings = outputs.last_hidden_state[:, 0]

        logits = self.classifier(img_embeddings)
        return logits

    def _unfreeze_last_n_layers(self, n: int):
        if hasattr(self.backbone, "encoder"):
            encoder_layers = self.backbone.encoder.layer
        elif hasattr(self.backbone, "vision_model"):
            encoder_layers = self.backbone.vision_model.encoder.layer
        else:
            raise ValueError("Cannot find encoder layers in backbone.")

        total_layers = len(encoder_layers)
        n = min(n, total_layers)

        print(f"Unfreezing last {n} of {total_layers} transformer layers.")

        for layer in encoder_layers[-n:]:
            for param in layer.parameters():
                param.requires_grad = True


# 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, backbone_model_name, num_layers_to_unfreeze, num_classes=configs['NUM_CLASSES']):
        super().__init__()
        self.num_classes = num_classes
        self.backbone_model_name = backbone_model_name
        self.num_layers_to_unfreeze = num_layers_to_unfreeze

        # Load a pretrained backbone and replace its final layer
        self.model = get_pretrained_model(
            num_classes = self.num_classes,
            model_name  = self.backbone_model_name,
            num_layers_to_unfreeze = self.num_layers_to_unfreeze)

        # Binary classification loss operating on raw logits
        self.loss_function      = torch.nn.BCEWithLogitsLoss()

        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_logits = self.forward(x)
        y_prob    = torch.sigmoid(y_logits)

        # Compute metrics (expects logits + labels)
        loss     = self.loss_function(y_logits, 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_prob, y)
        f1_score = self.f1_score_function(y_prob, y)
        auroc    = self.auroc_function(y_prob, y)
        mAP      = self.map_score_function(y_prob, y) # mean average precision

        return loss, y_logits, y, accuracy, f1_score, auroc, mAP

    def training_step(self, batch, batch_idx):
        # Run shared step
        loss, y_logits, 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_logits, 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_logits, 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.parameters(), lr=3e-5)
        return optimizer

"""### Create function for running inference (i.e., assistive medical diagnosis)"""

def generate_query(formatted_predictions):
    return f"""
      The predicted conditions and their corresponding probabilities are given by the following dictionary:

      {formatted_predictions}

      What features of the chest X-ray image support the predicted condition(s)?
      """

def predictionReportGenerator(vlm_model, img_pil, system_prompt, query_prompt):
    # image_ = Image.open(image_path).convert("RGB")
    image_ = img_pil.convert("RGB")

    messages = [
    {
        "role": "system",
        "content": [{"type": "text", "text": f"{system_prompt}"}]},
    {
        "role": "user",
        "content": [
            {"type": "image", "image": image_},
            {"type": "text", "text": f"{query_prompt}"}]}]

    output = vlm_model(text=messages, max_new_tokens=350)
    prediction_explanation = output[0]["generated_text"][-1]["content"]
    
    return prediction_explanation


    
@torch.inference_mode()
def run_diagnosis(
    backbone_name,
    vlm_name,
    input_image,
    threshold,
    preprocess_fn=None,
    Idx2labels=None):

    # Preprocess
    x = preprocess_fn(input_image).unsqueeze(0)

    # Resolve backbone
    # ckpt_path = os.path.join(CKPT_ROOT, MODEL_REGISTRY[backbone_name])
    ckpt_path = os.path.join(CKPT_ROOT, f"{backbone_name}.ckpt")

    if not os.path.exists(ckpt_path):
        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")

    # Load classification model (cache for speed)
    if backbone_name not in ViT_MODEL_CACHE:
      ViT_MODEL_CACHE[backbone_name] = modelModule.load_from_checkpoint(
        ckpt_path, backbone_model_name=backbone_name, num_layers_to_unfreeze = 2)
    model = ViT_MODEL_CACHE[backbone_name]

    model.eval()

    # device = 0 if torch.cuda.is_available() else -1

    # Forward
    logits = model(x)
    probs = torch.sigmoid(logits)[0].cpu().numpy()

    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]

    explanation_ = "No prediction was made."
    if predicted_classes:
        # Load model (cache for speed)
        if vlm_name not in VLM_MODEL_CACHE:
          VLM_MODEL_CACHE[vlm_name]     = pipeline(task  = "image-text-to-text",
                              model  = VLM_REGISTRY[vlm_name],
                              trust_remote_code = True)
        VLM_model = VLM_MODEL_CACHE[vlm_name]

        formatted_predictions = {label: output_probs[label] for label in predicted_classes}
        
        query_prompt = generate_query(formatted_predictions)
        explanation_ = predictionReportGenerator(vlm_model = VLM_model, img_pil = input_image, 
                                                 system_prompt = VLM_SYSTEM_PROMPT, query_prompt = query_prompt)

    return "\n".join(predicted_classes), explanation_, output_probs

"""### Gradio app"""

CKPT_ROOT = os.path.join(os.getcwd(), "Vision Transformers")

example_list_dir = os.path.join(os.getcwd(), "Curated test samples")
example_list_img_names = os.listdir(example_list_dir)

example_list = [
    [configs["DEFAULT_BACKBONE"], configs["DEFAULT_VLM"], os.path.join(example_list_dir, example_img)]
    for example_img in example_list_img_names[:8]
    if example_img.lower().endswith(".png")]


gradio_app = gradio.Interface(
    fn     = partial(run_diagnosis, preprocess_fn = preprocess_fxn, Idx2labels = labels_dict),

    inputs = [gradio.Dropdown(["CheXFormer-small", "ViT-base-16"], value="CheXFormer-small", label="Select Classification Model"),
              gradio.Dropdown(["MedGemma-4b", "MedMO", "Lingshu-7B", "Qwen3-VL-2B"], value="Lingshu-7B", label="Select Explanation 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.2, label = "Set Prediction Threshold")],

    outputs = [gradio.Textbox(label="Predicted Medical Condition(s)"),
               gradio.Textbox(label="Prediction Report"),
               gradio.Label(label="Predicted Probabilities", show_label=False)],

    examples       = example_list,
    cache_examples = False,
    title          = "ChestVision-PRO",
    description    = "Vision-Transformer solutions for assistive medical diagnosis with Vision-Language-based prediction justification",
    article        = "Author: C. Foli (02.2026) | Website: coming soon...")

gradio_app.launch()