threshold added

app.py

@@ -63,8 +63,9 @@ def nms(final_boxes, scores, IOU_threshold=0):
 
     return final_boxes[pick]
 
-def detect_obj(input_image):
+def detect_obj(input_image, obj_threshold, bb_threshold):
     try:
+        
         image = np.array(input_image)
         image = cv2.resize(image, (H, W))
         img = image
@@ -75,12 +76,9 @@ def detect_obj(input_image):
         output = m.run(['reshape'], {"input": image})
         output = np.squeeze(output, axis=0)
 
-        THRESH=.25
-
-
         object_positions = np.concatenate(
-                [np.stack(np.where(output[..., 0]>=THRESH), axis=-1),
-                 np.stack(np.where(output[..., 5]>=THRESH), axis=-1)], axis=0
+                [np.stack(np.where(output[..., 0]>=obj_threshold), axis=-1),
+                 np.stack(np.where(output[..., 5]>=obj_threshold), axis=-1)], axis=0
         )
 
         selected_output = []
@@ -93,7 +91,7 @@ def detect_obj(input_image):
 
         for i,pos in enumerate(object_positions):
             for j in range(2):
-                if selected_output[i][j*5]>THRESH:
+                if selected_output[i][j*5]>obj_threshold:
                     output_box = np.array(output[pos[0]][pos[1]][pos[2]][(j*5)+1:(j*5)+5], dtype=float)
 
                     x_centre = (np.array(pos[1], dtype=float) + output_box[0])*32
@@ -116,7 +114,7 @@ def detect_obj(input_image):
 
         final_boxes = np.array(final_boxes)
 
-        nms_output = nms(final_boxes, final_scores, 0.3)
+        nms_output = nms(final_boxes, final_scores, bb_threshold)
 
         for i in nms_output:
             cv2.rectangle(
@@ -140,17 +138,32 @@ def detect_obj(input_image):
         return input_image
 
 
+
 with gr.Blocks(title="YOLOS Object Detection - ClassCat", css=".gradio-container {background:lightyellow;}") as demo:
     gr.HTML('<h1>Yolo Object Detection</h1>')
     gr.HTML("<h4>supported objects are [aeroplane,bicycle,bird,boat,bottle,bus,car,cat,chair,cow,diningtable,dog,horse,motorbike,person,pottedplant,sheep,sofa,train,tvmonitor]</h4>")
+    gr.HTML("<br>")
     with gr.Row():
         input_image = gr.Image(label="Input image", type="pil")
         output_image = gr.Image(label="Output image", type="pil")
+    gr.HTML("<br>")
+    gr.HTML("<h4>object centre detection threshold means the object centre will be considered a new object if it's value is above threshold</h4>")
+    gr.HTML("<p>less means more objects</p>")
+    gr.HTML("<h4>bounding box threshold is IOU value threshold. If intersection/union area of two bounding boxes are greater than threshold value the one box will be suppressed</h4>")
+    gr.HTML("<p>more means more bounding boxes<p>")
+    gr.HTML("<br>")
+
+    obj_threshold = gr.Slider(0, 1.0, value=0.2, label=' object centre detection threshold')
+    gr.HTML("<br>")
+    bb_threshold = gr.Slider(0, 1.0, value=0.3, label=' bounding box draw threshold')
+    gr.HTML("<br>")
 
     send_btn = gr.Button("Detect")
+    gr.HTML("<br>")
     gr.Examples(['./samples/out_1.jpg'], inputs=input_image)
 
-    send_btn.click(fn=detect_obj, inputs=[input_image], outputs=[output_image])
+    send_btn.click(fn=detect_obj, inputs=[input_image, obj_threshold, bb_threshold], outputs=[output_image])
+
 
 
 demo.launch(debug=True)