confidence threshold

app.py

@@ -2,21 +2,21 @@ import numpy as np
 import torch
 import gradio as gr
 from infer import detections
+'''
 import os
 os.system("mkdir data")
 os.system("mkdir data/models")
 os.system("wget https://www.cs.cmu.edu/~walt/models/walt_people.pth -O data/models/walt_people.pth")
 os.system("wget https://www.cs.cmu.edu/~walt/models/walt_vehicle.pth -O data/models/walt_vehicle.pth")
 '''
-'''
-def walt_demo(input_img):
+def walt_demo(input_img, confidence_threshold):
     #detect_people = detections('configs/walt/walt_people.py', 'cuda:0', model_path='data/models/walt_people.pth')
     if torch.cuda.is_available() == False:
         device='cpu'
     else:
         device='cuda:0'
     #detect_people = detections('configs/walt/walt_people.py', device, model_path='data/models/walt_people.pth')
-    detect = detections('configs/walt/walt_vehicle.py', device, model_path='data/models/walt_vehicle.pth', threshold=0.75)
+    detect = detections('configs/walt/walt_vehicle.py', device, model_path='data/models/walt_vehicle.pth', threshold=confidence_threshold)
  
     count = 0
     #img = detect_people.run_on_image(input_img)
@@ -45,9 +45,9 @@ article="""
 """
 
 examples = [
-    'demo/images/img_1.jpg',
-    'demo/images/img_2.jpg',
-    'demo/images/img_4.png',
+    ['demo/images/img_1.jpg',0.8],
+    ['demo/images/img_2.jpg',0.8],
+    ['demo/images/img_4.png',0.85],
 ]
 
 '''
@@ -58,9 +58,15 @@ img=walt_demo(img)
 cv2.imwrite(filename.replace('/images/','/results/'),img)
 cv2.imwrite('check.png',img)
 '''
+confidence_threshold = gr.Slider(minimum=0.3,
+                                    maximum=1.0,
+                                    step=0.01,
+                                    value=1.0,
+                                    label="Amodal Detection Confidence Threshold")
+inputs = [gr.Image(), confidence_threshold]
 demo = gr.Interface(walt_demo, 
-        gr.Image(), 
-        "image",
+        outputs="image",
+        inputs=inputs, 
         article=article,
         title=title,
         enable_queue=True,