ViTPose-ONNX

Easy inference for ViTPose using ONNX

Requirements

pip install -r requirements.txt

As you can see in 'requirements.txt', it requires only 5 libraries below

matplotlib
numpy
onnxruntime-gpu
opencv-python
yacs

Usage

Install

git clone https://github.com/Pukei-Pukei/ViTPose-ONNX.git
cd ViTPose-ONNX
pip install -r requirements.txt

Run

Download vitpose-b-multi-coco.onnx and yolov6m.onnx, then put them in ViTPose-ONNX folder
Run the commands below to start inference

python run.py -img <path_to_image>

python run.py -vid <path_to_video>

python run.py -wc <webcam ID or URL>

python run.py -cfg <config path> -vid <path_to_video>

Example

python run.py -cfg configs/custom_config.py -vid dance.mp4 -s

'-s' for save option

Options

--yolov6-path, -yolo PATH       :   Path to YOLOv6 onnx file
--vitpose-path, -pose PATH      :   Path to ViTPose onnx file

--image-path, -img PATH         :   Image path 
--video-path, -vid PATH         :   Videos path 
--webcam, -wc PATH              :   Webcam id or webcam URL 

--no-background, -nobg          :   Background will be black screen
--no-bbox, -nobx                :   Don't draw bboxes
--no-skeleton, -nosk            :   Don't draw skeletons
--dynamic-drawing, -dd          :   Turn on dynamic drawing, keypoint 
                                    radius and skeleton width change 
                                    dynamically with bbox size
--result-scale, -rs SIZE        :   Set a coefficient to scale a size 
                                    of result, set None for not 
                                    processing

--save, -s                      :   Save drawing result
--save-prediction, -sp          :   Save the predictions(bbox, pose), 
                                    Numpy is needed to read the save 
                                    file

--conf-thres, -conf THRES       :   Set confidence threshold for 
                                    non-maximum suppression
--iou-thres, -iou THRES         :   Set IoU threshold for 
                                    non-maximum suppression
--max-detection, -max MAX       :   Set max detection for non-maximum 
                                    suppression
--key-conf-thres, -kconf THRES  :   Set keypoint confidence threshold
--no-pad                        :   Don't use additional padding

--cpu, -cpu                     :   Use cpu instead of gpu
--pose-batch-size, -pbs SIZE    :   Set pose batch size
--yolo-batch-size, -ybs SIZE    :   Set yolo batch size, 
                                    it works only in video

--config, -cfg                  :   Config path. use config for easy 
                                    usage of options. default config 
                                    path is 'configs/base_config.py'

Download ONNX file

Model	ONNX	Original Weight for PyTorch
ViTPose-B	GoogleDrive	Onedrive
YOLOv6-M	GoogleDrive	Download

If you want other versions, refer to Tutorial and get your own ONNX

Pukei-Pukei / ViTPose-ONNX