README.md

---
title: RTMO Checkpoint Tester
emoji: 👀
colorFrom: pink
colorTo: green
sdk: gradio
sdk_version: 5.27.0
app_file: app.py
pinned: false
license: apache-2.0
short_description: RTMO PyTorch Checkpoint Tester
---

# RTMO PyTorch Checkpoint Tester

This HuggingFace Space provides a real-time 2D multi-person pose estimation demo using the RTMO model from OpenMMLab, accelerated with ZeroGPU. It supports both image and video inputs.

## Features

- **Remote Checkpoint Selection**: Choose from multiple pre-trained variants (COCO, BODY7, CrowdPose, retrainable RTMO-s) via a dropdown.
- **Custom Checkpoint Upload**: Upload your own `.pth` file; the application auto-detects RTMO-t/s/m/l variants.
- **Image Input**: Upload images for single-frame pose estimation.
- **Video Input**: Upload video files (e.g., `.mp4`, `.mov`, `.avi`, `.mkv`, `.webm`) to perform pose estimation on video sequences and view annotated outputs.
- **Threshold Adjustment**: Fine-tune **Bounding Box Threshold** and **NMS Threshold** sliders to refine detections.
- **Example Images**: Three license-free images with people are included for quick testing via the **Examples** panel.
- **ZeroGPU Acceleration**: Utilizes the `@spaces.GPU()` decorator for GPU inference on HuggingFace Spaces.

## Usage

1. **Upload Image**: Drag-and-drop or select an image in the **Upload Image** component (or choose from **Examples**).
2. **Upload Video**: Drag-and-drop or select a video file in the **Upload Video** component.
3. **Select Remote Checkpoint**: Pick a preloaded variant from the dropdown menu.
4. **(Optional) Upload Your Own Checkpoint**: Provide a `.pth` file to override the remote selection; the model variant is detected automatically.
5. **Adjust Thresholds**: Set **Bounding Box Threshold** (`bbox_thr`) and **NMS Threshold** (`nms_thr`) to control confidence and suppression behavior.
6. **Run Inference**: Click **Run Inference**.
7. **View Results**:
   - For images, the annotated image will appear in the **Annotated Image** panel.
   - For videos, the annotated video will appear in the **Annotated Video** panel.
   The active checkpoint name will appear below.

## Remote Checkpoints

The following variants are available out of the box:

- `rtmo-s_8xb32-600e_coco`
- `rtmo-m_16xb16-600e_coco`
- `rtmo-l_16xb16-600e_coco`
- `rtmo-t_8xb32-600e_body7`
- `rtmo-s_8xb32-600e_body7`
- `rtmo-m_16xb16-600e_body7`
- `rtmo-l_16xb16-600e_body7`
- `rtmo-s_8xb32-700e_crowdpose`
- `rtmo-m_16xb16-700e_crowdpose`
- `rtmo-l_16xb16-700e_crowdpose`
- `rtmo-s_coco_retrainable` (from Hugging Face)

## Implementation Details

- **GPU Decorator**: `@spaces.GPU()` marks the `predict` function for GPU execution under ZeroGPU.
- **Inference API**: Leverages `MMPoseInferencer` from MMPose with `pose2d`, `pose2d_weights`, and category `[0]` for person detection.
- **Monkey-Patch**: Applies a regex patch to bypass `mmdet`’s MMCV version assertion for compatibility.
- **Variant Detection**: Inspects `backbone.stem.conv.conv.weight` channels in the checkpoint to select the correct RTMO variant.
- **Checkpoint Management**: Remote files are downloaded to `/tmp/{key}.pth` on demand; uploads use the provided local path.
- **Image & Video Support**: The `predict` function automatically handles both image and video inputs, saving annotated frames or video to `/tmp/vis` and displaying them in the UI.
- **Output**: Saves visualization images or videos to `/tmp/vis` and displays them in the UI panels.

## Files

- **app.py**: Main Gradio application script.
- **requirements.txt**: Python dependencies, including MMCV and MMPose.
- **README.md**: This documentation file.