feat: Control threshold

app.py

@@ -18,10 +18,10 @@ model = RFDETRBase(pretrain_weights="checkpoint_best_regular.pth")
 box_annotator = sv.BoxAnnotator()
 label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
 
-def detect_objects_and_recognize_logos(image):
+def detect_objects_and_recognize_logos(image, threshold):
     try:
-        # Run inference with RFDETR
-        detections = model.predict(image, threshold=0.2)
+        # Run inference with RFDETR using the provided threshold
+        detections = model.predict(image, threshold=threshold)
         
         # Initialize labels for detection and recognition frames
         detection_labels = []
@@ -101,14 +101,17 @@ def detect_objects_and_recognize_logos(image):
 # Create Gradio interface
 interface = gr.Interface(
     fn=detect_objects_and_recognize_logos,
-    inputs=gr.Image(type="pil", label="Upload Image"),
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="Confidence Threshold")
+    ],
     outputs=[
         gr.Image(type="pil", label="Detection Frame (RFDETR)"),
         gr.Image(type="pil", label="Recognition Frame (RFDETR + Gemini)"),
         gr.Textbox(label="Detected Brand Names")
     ],
     title="Object Detection and Logo Recognition with RFDETR and Gemini",
-    description="Upload an image to detect objects using RFDETR model and recognize logos using Google Gemini. Outputs include a detection frame (objects only) and a recognition frame (objects with brand names)."
+    description="Upload an image to detect objects using RFDETR model and recognize logos using Google Gemini. Adjust the confidence threshold to filter detections. Outputs include a detection frame (objects only) and a recognition frame (objects with brand names)."
 )
 
 # Launch the interface