Update app.py

app.py

@@ -117,13 +117,14 @@ def get_model_info(model_info):
     )
     return info_text
 
-def predict_image(model_name, image, models):
+def predict_image(model_name, image, confidence, models):
     """
     Perform prediction on an uploaded image using the selected YOLO model.
     
     Args:
         model_name (str): The name of the selected model.
         image (PIL.Image.Image): The uploaded image.
+        confidence (float): The confidence threshold for detections.
         models (dict): The dictionary containing models and their info.
     
     Returns:
@@ -142,8 +143,8 @@ def predict_image(model_name, image, models):
         input_image_path = os.path.join(TEMP_DIR, f"{model_name}_input_image.jpg")
         image.save(input_image_path)
         
-        # Perform prediction
-        results = model(input_image_path, save=True, save_txt=False, conf=0.25)
+        # Perform prediction with user-specified confidence
+        results = model(input_image_path, save=True, save_txt=False, conf=confidence)
         
         # Determine the output path
         # Ultralytics YOLO saves the results in 'runs/detect/predict' by default
@@ -165,13 +166,14 @@ def predict_image(model_name, image, models):
     except Exception as e:
         return f"❌ Error during prediction: {str(e)}", None, None
 
-def predict_video(model_name, video, models):
+def predict_video(model_name, video, confidence, models):
     """
     Perform prediction on an uploaded video using the selected YOLO model.
     
     Args:
         model_name (str): The name of the selected model.
         video (str): Path to the uploaded video file.
+        confidence (float): The confidence threshold for detections.
         models (dict): The dictionary containing models and their info.
     
     Returns:
@@ -190,8 +192,8 @@ def predict_video(model_name, video, models):
         input_video_path = os.path.join(TEMP_DIR, f"{model_name}_input_video.mp4")
         shutil.copy(video, input_video_path)
         
-        # Perform prediction
-        results = model(input_video_path, save=True, save_txt=False, conf=0.25)
+        # Perform prediction with user-specified confidence
+        results = model(input_video_path, save=True, save_txt=False, conf=confidence)
         
         # Determine the output path
         # Ultralytics YOLO saves the results in 'runs/detect/predict' by default
@@ -253,6 +255,17 @@ def main():
             outputs=model_info
         )
         
+        # Confidence Threshold Slider
+        with gr.Row():
+            confidence_slider = gr.Slider(
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.25,
+                label="Confidence Threshold",
+                info="Adjust the minimum confidence required for detections to be displayed."
+            )
+        
         # Tabs for different input types
         with gr.Tabs():
             # Image Prediction Tab
@@ -269,16 +282,16 @@ def main():
                     image_download_btn = gr.File(label="⬇️ Download Predicted Image")
                 
                 # Define the image prediction function
-                def process_image(selected_display_name, image):
+                def process_image(selected_display_name, image, confidence):
                     if not selected_display_name:
                         return "❌ Please select a model.", None, None
                     model_name = display_to_name.get(selected_display_name)
-                    return predict_image(model_name, image, models)
+                    return predict_image(model_name, image, confidence, models)
                 
                 # Connect the predict button
                 image_predict_btn.click(
                     fn=process_image,
-                    inputs=[model_dropdown, image_input],
+                    inputs=[model_dropdown, image_input, confidence_slider],
                     outputs=[image_status, image_output, image_download_btn]
                 )
             
@@ -294,16 +307,16 @@ def main():
                     video_download_btn = gr.File(label="⬇️ Download Predicted Video")
                 
                 # Define the video prediction function
-                def process_video(selected_display_name, video):
+                def process_video(selected_display_name, video, confidence):
                     if not selected_display_name:
                         return "❌ Please select a model.", None, None
                     model_name = display_to_name.get(selected_display_name)
-                    return predict_video(model_name, video, models)
+                    return predict_video(model_name, video, confidence, models)
                 
                 # Connect the predict button
                 video_predict_btn.click(
                     fn=process_video,
-                    inputs=[model_dropdown, video_input],
+                    inputs=[model_dropdown, video_input, confidence_slider],
                     outputs=[video_status, video_output, video_download_btn]
                 )