Update app.py

app.py

@@ -7,7 +7,7 @@ from PIL import Image, ImageDraw
 processor = AutoImageProcessor.from_pretrained("0llheaven/Conditional-detr-finetuned")
 model = AutoModelForObjectDetection.from_pretrained("0llheaven/Conditional-detr-finetuned")
 
-def detect_objects(image):
+def detect_objects(image, score_threshold):
     # Convert image to RGB if it's grayscale
     if image.mode != "RGB":
         image = image.convert("RGB")
@@ -16,7 +16,7 @@ def detect_objects(image):
     inputs = processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
 
-    # Filter predictions with confidence greater than 0.5
+    # Filter predictions based on the user-defined score threshold
     target_sizes = torch.tensor([image.size[::-1]])
     results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)
 
@@ -28,17 +28,18 @@ def detect_objects(image):
         boxes = result["boxes"]
 
         for score, label, box in zip(scores, labels, boxes):
-            box = [round(i, 2) for i in box.tolist()]
-            label_name = "Pneumonia" if label.item() == 0 else "Other"
-            draw.rectangle(box, outline="red", width=3)
-            draw.text((box[0], box[1]), f"{label_name}: {round(score.item(), 3)}", fill="red")
+            if score >= score_threshold:  # Only draw if score is above threshold
+                box = [round(i, 2) for i in box.tolist()]
+                label_name = "Pneumonia" if label.item() == 0 else "Other"
+                draw.rectangle(box, outline="red", width=3)
+                draw.text((box[0], box[1]), f"{label_name}: {round(score.item(), 3)}", fill="red")
     
     return image
 
 # Create the Gradio interface
 interface = gr.Interface(
     fn=detect_objects, 
-    inputs=gr.Image(type="pil"), 
+    inputs=[gr.Image(type="pil"), gr.Slider(0, 1, value=0.5, label="Score Threshold")],  # Add slider for score threshold
     outputs=gr.Image(type="pil"),  # Corrected output type
     title="Object Detection with Transformers",
     description="Upload an image to detect objects using a fine-tuned Conditional-DETR model."