Let user configurate bbox_thr & nms_thr in GUI

README.md

@@ -28,4 +28,9 @@ This HuggingFace Space runs the RTMO (Real-Time Multi-Person) 2D pose estimation
 We use the `rtmo` alias defined in MMPose’s model zoo. To override, upload your own checkpoint.
 
 ## Development
-If you need to update dependencies or change the model, modify `requirements.txt` and `app.py` accordingly.
+If you need to update dependencies or change the model, modify `requirements.txt` and `app.py` accordingly.
+
+## Todos
+1. Let user configurate bbox_thr & nms_thr in GUI
+2. Support video input
+3. Support models in ONNX format via rtmlib

app.py

@@ -56,7 +56,7 @@ def detect_rtmo_variant(checkpoint_path: str) -> str:
 
 # ——— Gradio prediction function ———
 @spaces.GPU()
-def predict(image: Image.Image, checkpoint):
+def predict(image: Image.Image, checkpoint, bbox_thr: float, nms_thr: float):
     # save upload to temp file
     inp_path = "/tmp/upload.jpg"
     image.save(inp_path)
@@ -70,8 +70,8 @@ def predict(image: Image.Image, checkpoint):
     inferencer = load_inferencer(checkpoint_path=ckpt_path, device=None)
     for result in inferencer(
         inputs=inp_path,
-        bbox_thr=0.1,
-        nms_thr=0.65,
+        bbox_thr=bbox_thr,
+        nms_thr=nms_thr,
         pose_based_nms=True,
         show=False,
         vis_out_dir=vis_dir,
@@ -89,6 +89,8 @@ demo = gr.Interface(
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
         gr.File(file_types=['.pth'], label="Upload RTMO .pth Checkpoint (optional)")
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.1, label="Bounding Box Threshold (bbox_thr)"),
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.65, label="NMS Threshold (nms_thr)"),
     ],
     outputs=gr.Image(type="pil", label="Annotated Image"),
     title="RTMO Pose Demo",