Update app.py

app.py

@@ -80,7 +80,7 @@ with gr.Blocks() as demo:
         
 
     # main function of the app
-    def yolo3_inference(input_img,gradcam=True,gradcam_opa=0.5): # function for yolo inference
+    def yolo3_inference(input_img,gradcam=True,gradcam_opa=0.5,user_iou_threshold=0.6,user_threshold=0.5): # function for yolo inference
       
       # load model
       yololit = litmodelclass.LitYolo()
@@ -127,7 +127,7 @@ with gr.Blocks() as demo:
           for idx, (box) in enumerate(boxes_scale_i):
               bboxes[idx] += box
       # nms
-      nms_boxes = utils.non_max_suppression(bboxes[0], iou_threshold=0.6, threshold=0.5, box_format="midpoint",)
+      nms_boxes = utils.non_max_suppression(bboxes[0], iou_threshold=user_iou_threshold, threshold=user_threshold, box_format="midpoint",)
       nms_boxes_output.append(nms_boxes)
 
       # use gradio image annotations
@@ -173,13 +173,15 @@ with gr.Blocks() as demo:
     with gr.Row():
       gradcam_check = gr.Checkbox(label="Gradcam")
       gradcam_opa = gr.Slider(0, 1, value = 0.5, label="Opacity of GradCAM")
+      iou_threshold = gr.Slider(0, 1, value = 0.6, label="IOU Threshold")
+      threshold = gr.Slider(0, 1, value = 0.5, label="Threshold")
         
     section_btn = gr.Button("Identify Objects")
-    section_btn.click(yolo3_inference, inputs=[img_input,gradcam_check,gradcam_opa], outputs=[img_output])
+    section_btn.click(yolo3_inference, inputs=[img_input,gradcam_check,gradcam_opa,iou_threshold,threshold], outputs=[img_output])
 
     gr.Markdown("## Some Examples")
     gr.Examples(examples=examples,
-                             inputs =[img_input,gradcam_check,gradcam_opa],
+                             inputs =[img_input,gradcam_check,gradcam_opa,iou_threshold,threshold],
                              outputs=img_output,
                              fn=yolo3_inference, cache_examples=False)