Update app.py

app.py

@@ -7,28 +7,27 @@ import os
 def load_model(repo_id):
     download_dir = snapshot_download(repo_id)
     print(download_dir)
-    path  = os.path.join(download_dir, "best_int8_openvino_model")
+    path = os.path.join(download_dir, "best_int8_openvino_model")
     print(path)
     detection_model = YOLO(path, task='detect')
     return detection_model
 
-
-def predict(pilimg):
-
+def predict(pilimg, conf_threshold, iou_threshold):
     source = pilimg
-    # x = np.asarray(pilimg)
-    # print(x.shape)
-    result = detection_model.predict(source, conf=0.5, iou=0.6)
+    result = detection_model.predict(source, conf=conf_threshold, iou=iou_threshold)
     img_bgr = result[0].plot()
-    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
-    
+    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB for PIL
     return out_pilimg
 
-
 REPO_ID = "sensura/belisha-beacon-zebra-crossing-yoloV8"
 detection_model = load_model(REPO_ID)
 
-gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil"),
-             outputs=gr.Image(type="pil")
-             ).launch(share=True)
+gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="Confidence Threshold"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.6, step=0.05, label="IoU Threshold")
+    ],
+    outputs=gr.Image(type="pil", label="Detection Output")
+).launch(share=True)