Update app.py

app.py

@@ -21,6 +21,7 @@ def detect_objects(image, score_threshold):
     results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)
 
     labels_output = [] 
+    no_detection = True
     
     # Draw bounding boxes around detected objects
     draw = ImageDraw.Draw(image)
@@ -31,11 +32,26 @@ def detect_objects(image, score_threshold):
 
         for score, label, box in zip(scores, labels, boxes):
             if score >= score_threshold:  # Only draw if score is above threshold
+                no_detection = False
                 box = [round(i, 2) for i in box.tolist()]
-                label_name = "Pneumonia" if label.item() == 0 else "Other"
+                
+                #label_name = "Pneumonia" if label.item() == 0 else "Other"   
+                if label.item() == 0:
+                    label_name = "Pneumonia"
+                elif label.item() == 1:
+                    label_name = "Normal"
+                elif label.item() == 2:
+                    label_name = "Pneumonia_bacteria"
+                else:
+                    label_name = "Pneumonia_virus"
+                    
                 draw.rectangle(box, outline="red", width=3)
                 draw.text((box[0], box[1]), f"{label_name}: {round(score.item(), 3)}", fill="red")
                 labels_output.append(f"{label_name}: {round(score.item(), 3)}")
+
+    # If no detections, set label as 'Other'
+    if no_detection:
+        labels_output.append("Other")
     
     return image, "\n".join(labels_output)