Update app.py

app.py

@@ -46,6 +46,25 @@ configs = {
     "THRESHOLD": 0.5
     }
 
+BACKBONE_REGISTRY = {
+    "EfficientNet(b3)": {
+        "torchvision_name": "efficientnet_b3",
+        "ckpt": "EfficientNet(b3).ckpt"},
+    "ConvNeXt(tiny)": {
+        "torchvision_name": "convnext_tiny",
+        "ckpt": "ConvNeXt(tiny).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
@@ -211,39 +230,46 @@ class modelModule(torch_light.LightningModule):
 
 """### Create function for running inference (i.e., assistive medical diagnosis)"""
 
-def run_diagnosis(backbone_name, input_image, preprocess_fn = None, Idx2labels=None, threshold = configs["THRESHOLD"]):
+@torch.inference_mode()
+def run_diagnosis(
+    backbone_name,
+    input_image,
+    preprocess_fn=None,
+    Idx2labels=None,
+    threshold=configs["THRESHOLD"]):
+  
+    # 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:
+      MODEL_CACHE[backbone_name] = modelModule.load_from_checkpoint(
+        ckpt_path, backbone_model_name=backbone_info["torchvision_name"])
+    model = MODEL_CACHE[backbone_name]
 
-    input_tensor = preprocess_fn(input_image)
-    input_tensor = input_tensor.unsqueeze(dim = 0)
-    # newimg = transform(img).unsqueeze(dim=0)
-
-    CKPT_PATH = os.path.join(CKPT_ROOT, f"{backbone_name}.ckpt")
-    model = modelModule.load_from_checkpoint(CKPT_PATH)
     model.eval()
 
-    # Generate predictions
-    output_logits = model(input_tensor).cpu()
-
-    # File "/app/app.py", line 226, in run_diagnosis
-    # probabilities = torch.sigmoid(output_logits)[0].numpy().tolist()
-    #                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
-    # RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.
+    # Forward
+    logits = model(x)
+    probs = torch.sigmoid(logits)[0].cpu().numpy()
 
-    #    ylogit = model(newimg).detach()
-    #     yprob = torch.sigmoid(ylogit)
-
-    probabilities = torch.sigmoid(output_logits)[0].detach().numpy().tolist()
+    output_probs = {
+        Idx2labels[i]: float(p) for i, p in enumerate(probs)
+    }
 
-    output_probs = dict()
-    predicted_classes = []
+    predicted_classes = [
+        Idx2labels[i] for i, p in enumerate(probs) if p >= threshold
+    ]
 
-    for idx, prob in enumerate(probabilities):
-        output_probs[Idx2labels[idx]] = prob
-        if prob >= threshold:
-            predicted_classes.append(Idx2labels[idx])
+    return "\n".join(predicted_classes), output_probs
 
-    predicted_classes = "\n".join(predicted_classes)
-    return predicted_classes, output_probs
 
 """### Gradio app"""
 CKPT_ROOT = os.path.join(os.getcwd(), "Trained models")