Update app.py

app.py

@@ -144,29 +144,39 @@ def infer_mask_and_mask_image(pil_img, threshold=0.5):
 def classify_masked_tensor(masked_img_tensor, thresh_b=0.5, thresh_v=0.5):
     """
     masked_img_tensor: C,H,W on device, normalized for classifier
-    returns (pb, pv, label)
-    pb = probability of pneumonia class from model_bact (index 1)
-    pv = probability of pneumonia class from model_viral (index 1)
+    Returns (pb, pv, label)
+    pb = probability pneumonia in bacterial model
+    pv = probability pneumonia in viral model
     """
     x = masked_img_tensor.unsqueeze(0).to(device)
+
     with torch.no_grad():
-        out_b = model_bact(x)  
+        out_b = model_bact(x)
         out_v = model_viral(x)
-        prob_b = torch.softmax(out_b, dim=1)[0,1].item()
-        prob_v = torch.softmax(out_v, dim=1)[0,1].item()
 
-    # Decision logic: thresholds + fallback to higher prob when both triggered
-    if prob_b < thresh_b and prob_v < thresh_v:
+        pb = torch.softmax(out_b, dim=1)[0,1].item()  
+        pv = torch.softmax(out_v, dim=1)[0,1].item()  
+
+    # ----------- DECISION LOGIC -----------
+    # Case 1: Both low → NORMAL
+    if pb < thresh_b and pv < thresh_v:
         label = "NORMAL"
-    elif prob_b >= thresh_b and prob_v < thresh_v:
+
+    # Case 2: Only bacterial high → BACTERIAL
+    elif pb >= thresh_b and pv < thresh_v:
         label = "BACTERIAL PNEUMONIA"
-    elif prob_v >= thresh_v and prob_b < thresh_b:
+
+    # Case 3: Only viral high → VIRAL
+    elif pv >= thresh_v and pb < thresh_b:
         label = "VIRAL PNEUMONIA"
+
+    # Case 4: Both high → pick the dominant type
     else:
-        # both triggered -> pick the stronger probability (fallback)
-        label = "BACTERIAL PNEUMONIA" if prob_b > prob_v else "VIRAL PNEUMONIA"
-        label = label + " (fallback)"
-    return prob_b, prob_v, label
+        label = "BACTERIAL PNEUMONIA" if pb > pv else "VIRAL PNEUMONIA"
+        label += " (fallback-high-confidence-overlap)"
+
+    return pb, pv, label
+
 
 
 def inference_pipeline(img, thresh_b=0.5, thresh_v=0.5, seg_thresh=0.5):