Update app.py

app.py

@@ -7,10 +7,6 @@ Original file is located at
     https://colab.research.google.com/drive/1gVrx5TyipNPvn8D7GaK0pNBCnLeYTAD_
 """
 
-!pip install --upgrade gradio
-
-!pip install lightning torchmetrics
-
 """### Import dependencies"""
 
 import torch
@@ -256,8 +252,8 @@ def run_diagnosis(
     x = preprocess_fn(input_image).unsqueeze(0)
 
     # Resolve backbone
-    backbone_info = MODEL_REGISTRY[backbone_name]
-    ckpt_path = os.path.join(CKPT_ROOT, backbone_info["ckpt"])
+    # backbone_info = MODEL_REGISTRY[backbone_name]
+    ckpt_path = os.path.join(CKPT_ROOT, MODEL_REGISTRY[backbone_name])
 
     if not os.path.exists(ckpt_path):
         raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
@@ -265,7 +261,7 @@ def run_diagnosis(
     # Load model (cache for speed)
     if backbone_name not in MODEL_CACHE:
       MODEL_CACHE[backbone_name] = modelModule.load_from_checkpoint(
-        ckpt_path, backbone_model_name=backbone_info["torchvision_name"], num_layers_to_unfreeze = 2)
+        ckpt_path, backbone_model_name=MODEL_REGISTRY[backbone_name], num_layers_to_unfreeze = 2)
     model = MODEL_CACHE[backbone_name]
 
     model.eval()
@@ -298,19 +294,18 @@ example_list = [
 # example_list = [['/content/new_labels.csv',"ResNet50"]]
 
 gradio_app = gradio.Interface(
-    fn     = partial(run_diagnosis, preprocess_fn = preprocess_fxn, Idx2labels = labels_dict, threshold = configs["THRESHOLD"]),
+    fn     = partial(run_diagnosis, preprocess_fn = preprocess_fxn, Idx2labels = labels_dict),
 
-    # inputs = [gradio.Dropdown(["ConvNeXt(small)", "ConvNeXt(tiny)", "EfficientNet(v2_small)", "EfficientNet(b3)", "RegNet(x3_2GF)","ResNet50"], value="EfficientNet(b3)", label="Select Backbone Model"),
-    #           gradio.Image(type="pil", label="Load chest-X-ray image here")],
     inputs = [gradio.Dropdown(["CheXFormer-small", "ViT-base-16"], value="ViT-base-16", label="Select Backbone Model"),
-              gradio.Image(type="pil", label="Load chest-X-ray image here")],
+              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 Conditions"),
+    outputs = [gradio.Textbox(label="Predicted Medical Condition(s)"),
              gradio.Label(label="Predicted Probabilities", show_label=False)],
 
     examples       = example_list,
     cache_examples = True,
-    title          = "ChestVision",
+    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...")