Update app.py

app.py

@@ -45,6 +45,7 @@ def annotate_image(
 @spaces.GPU
 def process_video(
     input_video,
+    confidence_threshold,
     progress=gr.Progress(track_tqdm=True)
 ):
     video_info = sv.VideoInfo.from_video_path(input_video)
@@ -59,7 +60,7 @@ def process_video(
     with sv.VideoSink(result_file_path, video_info=video_info) as sink:
         for _ in tqdm(range(total), desc="Processing video.."):
             frame = next(frame_generator)
-            results = query(Image.fromarray(frame))
+            results = query(Image.fromarray(frame), confidence_threshold)
             final_labels = []
             detections = []
             
@@ -76,13 +77,13 @@ def process_video(
     return result_file_path
     
 
-def query(image):
+def query(image, confidence_threshold):
   inputs = processor(images=image, return_tensors="pt").to(device)
   with torch.no_grad():
     outputs = model(**inputs)
   target_sizes = torch.Tensor([image.size])
   
-  results = processor.post_process_object_detection(outputs=outputs, threshold=0.6, target_sizes=target_sizes)
+  results = processor.post_process_object_detection(outputs=outputs, threshold=confidence_threshold, target_sizes=target_sizes)
   return results
 
 with gr.Blocks() as demo:
@@ -94,6 +95,7 @@ with gr.Blocks() as demo:
         input_video = gr.Video(
             label='Input Video'
         )
+        conf = gr.Slider(label="Confidence Threshold", minimum=0.1, maximum=1.0, value=0.6, step=0.05)
         submit = gr.Button()
     with gr.Column():
         output_video = gr.Video(
@@ -101,16 +103,17 @@ with gr.Blocks() as demo:
         )
   gr.Examples(
       fn=process_video,
-      examples=[["./cats.mp4"]],
+      examples=[["./cats.mp4", 0.6]],
       inputs=[
-          input_video
+          input_video,
+          conf
       ],
       outputs=output_video
   )
 
   submit.click(
       fn=process_video,
-      inputs=input_video,
+      inputs=[input_video, conf],
       outputs=output_video
   )