add gradio ui

README.md

@@ -11,4 +11,21 @@ license: apache-2.0
 short_description: RTMO PyTorch Checkpoint Tester
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# RTMO PyTorch Checkpoint Tester
+
+This HuggingFace Space runs the RTMO (Real-Time Multi-Person) 2D pose estimation model from OpenMMLab.
+
+## Usage
+1. Upload an image via the Gradio UI.
+2. (Optional) Provide a path or URL to your own RTMO PyTorch `.pth` checkpoint. If left blank, the default pretrained weights will be used.
+3. Click **Submit**. The annotated image with keypoints will be displayed.
+
+## Files
+- **app.py**: Gradio application script that loads the RTMO model, runs inference, and displays results.
+- **requirements.txt**: Python dependencies, including the patched MMCV build and MMPose.
+
+## Model
+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.

app.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+import spaces
+import os, sys, importlib.util, re
+import gradio as gr
+from PIL import Image
+
+# ——— Monkey-patch mmdet to remove its mmcv-version assertion ———
+spec = importlib.util.find_spec('mmdet')
+if spec and spec.origin:
+    src = open(spec.origin, encoding='utf-8').read()
+    # strip out the mmcv_minimum_version…assert… block up to __all__
+    patched = re.sub(r'(?ms)^[ \t]*mmcv_minimum_version.*?^__all__', '__all__', src)
+    m = importlib.util.module_from_spec(spec)
+    m.__loader__ = spec.loader
+    m.__file__ = spec.origin
+    m.__path__ = spec.submodule_search_locations
+    sys.modules['mmdet'] = m
+    exec(compile(patched, spec.origin, 'exec'), m.__dict__)
+
+from mmpose.apis.inferencers import MMPoseInferencer
+
+# ——— Initialize inferencer with default RTMO 2D model ———
+def load_inferencer(checkpoint_path=None, device=None):
+    kwargs = {'pose2d': 'rtmo', 'scope': 'mmpose', 'device': device, 'det_cat_ids': [0]}
+    if checkpoint_path:
+        kwargs['pose2d_weights'] = checkpoint_path
+    return MMPoseInferencer(**kwargs)
+
+# ——— Gradio prediction function ———
+@spaces.GPU()
+def predict(image: Image.Image, checkpoint: str = None):
+    # save upload to temp file
+    inp_path = "/tmp/upload.jpg"
+    image.save(inp_path)
+
+    vis_dir = "/tmp/vis"
+    os.makedirs(vis_dir, exist_ok=True)
+
+    inferencer = load_inferencer(checkpoint_path=checkpoint, device=None)
+    # run inference & visualization
+    for result in inferencer(
+        inputs=inp_path,
+        bbox_thr=0.1,
+        nms_thr=0.65,
+        pose_based_nms=True,
+        show=False,
+        vis_out_dir=vis_dir,
+    ):
+        pass
+
+    # return the first visualization
+    out_files = sorted(os.listdir(vis_dir))
+    if out_files:
+        return Image.open(os.path.join(vis_dir, out_files[0]))
+    return None
+
+# ——— Build Gradio Interface ———
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.inputs.Image(type="pil", label="Upload Image"),
+        gr.inputs.Text(label="RTMO PyTorch Checkpoint Path (optional)")
+    ],
+    outputs=gr.outputs.Image(type="pil", label="Annotated Image"),
+    title="RTMO Pose Demo",
+    description="Upload an image, optionally supply a RTMO .pth checkpoint, and see 2D pose annotation.",
+)
+
+def main():
+    demo.launch()
+
+if __name__ == "__main__":
+    main()