Update app.py

app.py

@@ -30,7 +30,7 @@ import gradio
 from functools import partial
 from transformers import AutoModel
 
-"""### Set parameters"""
+"""### Initialize Containers"""
 
 configs = {
     "IMAGE_SIZE":   (224, 224),    # Resize images to (W, H)
@@ -46,13 +46,40 @@ configs = {
     "THRESHOLD": 0.2
     }
 
-MODEL_REGISTRY = {
+ViT_REGISTRY = {
     "CheXFormer-small": "m42-health/CXformer-small",
     # "CheXFormer-base": "m42-health/CXformer-base",
-    "ViT-base-16": "google/vit-base-patch16-224",
-}
+    "ViT-base-16": "google/vit-base-patch16-224"}
 
-MODEL_CACHE = {}
+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"""
 
@@ -69,7 +96,7 @@ class get_pretrained_model(nn.Module):
         print(f"Loading pretrained [{model_name}] model")
 
         self.backbone = AutoModel.from_pretrained(
-            MODEL_REGISTRY[model_name],
+            ViT_REGISTRY[model_name],
             # model_name,
             trust_remote_code=True)
 
@@ -243,9 +270,38 @@ class modelModule(torch_light.LightningModule):
 
 """### 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, image_path, system_prompt, query_prompt):
+    image_ = Image.open(image_path).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,
@@ -261,14 +317,16 @@ def run_diagnosis(
     if not os.path.exists(ckpt_path):
         raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
 
-    # Load model (cache for speed)
+    # Load classification model (cache for speed)
     if backbone_name not in MODEL_CACHE:
-      MODEL_CACHE[backbone_name] = modelModule.load_from_checkpoint(
+      ViT_MODEL_CACHE[backbone_name] = modelModule.load_from_checkpoint(
         ckpt_path, backbone_model_name=backbone_name, num_layers_to_unfreeze = 2)
-    model = MODEL_CACHE[backbone_name]
+    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()
@@ -278,10 +336,24 @@ def run_diagnosis(
     }
 
     predicted_classes = [
-        Idx2labels[i] for i, p in enumerate(probs) if p >= threshold
-    ]
+        Idx2labels[i] for i, p in enumerate(probs) if p >= threshold]
+
+    explanation_ = "No prediction was made."
+    if predicted_classes not []:
+        # Load model (cache for speed)
+        if model_key not in MODEL_CACHE:
+          VLM_MODEL_CACHE[model_key]     = pipeline(task  = "image-text-to-text",
+                              model  = VLM_REGISTRY[vlm_name],
+                              trust_remote_code = True)
+        VLM_model = VLM_MODEL_CACHE[model_key]
+
+        formatted_predictions = {label: output_probs[label] for label in predicted_classes}
+        
+        query_prompt = generate_query(formatted_predictions)
+        explanation_ = predictionReportGenerator(vlm_model = VLM_model, image_path = input_image, 
+                                                 system_prompt = VLM_SYSTEM_PROMPT, query_prompt = query_prompt)
 
-    return "\n".join(predicted_classes), output_probs
+    return "\n".join(predicted_classes), explanation_, output_probs
 
 """### Gradio app"""
 
@@ -299,12 +371,14 @@ example_list = [
 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 Backbone Model"),
+    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.Label(label="Predicted Probabilities", show_label=False)],
+               gradio.Textbox(label="Prediction Report"),
+               gradio.Label(label="Predicted Probabilities", show_label=False)],
 
     examples       = example_list,
     cache_examples = False,