Confidence Threshold

app.py

@@ -8,12 +8,13 @@ model_names = ["Nano", "Smoll", "Medium"]
 models = [YOLO(os.path.join("./saved_model", option)) for option in model_options]
 example_list = [["examples/" + example] for example in os.listdir("examples")]
 
-def process_image(input_image, model_name):
+def process_image(input_image, model_name, conf):
     print(model_name)
+    print(conf)
     model_index = model_names.index(model_name)
     model = models[model_index]
     
-    results = model.predict(input_image, conf=0.6)
+    results = model.predict(input_image, conf=conf)
     class_counts = {}
     class_counts_str = "Class Counts:\n"
 
@@ -32,7 +33,11 @@ def process_image(input_image, model_name):
 
 iface = gr.Interface(
     fn=process_image,
-    inputs=[gr.Image(), gr.Radio(model_names, label="Choose model", value=model_names[0])], 
+    inputs=[
+        gr.Image(), 
+        gr.Radio(model_names, label="Choose model", value=model_names[0]),
+        gr.Slider(minimum=0.2, maximum=1.0, step=0.1, label="Confidence Threshold", value=0.6)
+    ], 
     outputs=["image", gr.Textbox(label="More info")],
     title="YOLO Object detection. Trained on xView dataset.",
     description='''The xView dataset is composed of satellite images collected from WorldView-3 satellites at a 0.3m ground sample distance.\n 
@@ -40,6 +45,6 @@ iface = gr.Interface(
                    https://challenge.xviewdataset.org ''',
     live=True,
     examples=example_list
-    )
+)
 
 iface.launch()