YOLOv8n-Pose inherits the powerful object detection backbone and neck architecture of YOLOv8n. It extends the standard YOLOv8n object detection model by integrating dedicated pose estimation layers onto its head. This allows it to not only detect people (bboxes) but also simultaneously predict the spatial positions (keypoints) of their anatomical joints (e.g., shoulders, elbows, knees, ankles).

Inference results on Edge2.

Inference speed test: USB camera about 52ms per frame. MIPI camera about 40ms per frame.

Download YOLOv8 official code ultralytics/ultralytics

$ git clone https://github.com/ultralytics/ultralytics.git

Refer README.md to train a YOLOv8n-Pose model. My version torch==1.10.1 and ultralytics==8.0.86.

The SDK only supports python3.6 or python3.8, here is an example of creating a virtual environment for python3.8.

Install python packages.

$ sudo apt update
$ sudo apt install python3-dev python3-numpy

Follow this docs to install conda.

Then create a virtual environment.

$ conda create -n npu-env python=3.8
$ conda activate npu-env     #activate
$ conda deactivate           #deactivate

Download Tool from rockchip-linux/rknn-toolkit2.

$ git clone https://github.com/rockchip-linux/rknn-toolkit2.git
$ git checkout 9ad79343fae625f4910242e370035fcbc40cc31a

Install dependences and RKNN toolkit2 packages,

$ cd rknn-toolkit2
$ sudo apt-get install python3 python3-dev python3-pip
$ sudo apt-get install libxslt1-dev zlib1g-dev libglib2.0 libsm6 libgl1-mesa-glx libprotobuf-dev gcc cmake
$ pip3 install -r doc/requirements_cp38-*.txt
$ pip3 install packages/rknn_toolkit2-*-cp38-cp38-linux_x86_64.whl

After training model, modify Class Detect and Class Pose in ultralytics/ultralytics/nn/modules/head.py as follows. (If you use ultralytics==8.0.86, the class in ultralytics/ultralytics/nn/modules.py)

head.py

diff --git a/ultralytics/nn/modules/head.py b/ultralytics/nn/modules/head.py
index 0b02eb3..0a6e43a 100644
--- a/ultralytics/nn/modules/head.py
+++ b/ultralytics/nn/modules/head.py
@@ -42,6 +42,9 @@ class Detect(nn.Module):
 
     def forward(self, x):
         """Concatenates and returns predicted bounding boxes and class probabilities."""
+        if torch.onnx.is_in_onnx_export():
+            return self.forward_export(x)
+
         shape = x[0].shape  # BCHW
         for i in range(self.nl):
             x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
@@ -80,6 +83,15 @@ class Detect(nn.Module):
             a[-1].bias.data[:] = 1.0  # box
             b[-1].bias.data[:m.nc] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)
 
+    def forward_export(self, x):
+        results = []
+        for i in range(self.nl):
+            dfl = self.cv2[i](x[i]).contiguous()
+            cls = self.cv3[i](x[i]).contiguous()
+            results.append(torch.cat([cls, dfl], 1).permute(0, 2, 3, 1).unsqueeze(1))
+            # results.append(torch.cat([cls, dfl], 1))
+        return tuple(results)
+
 
@@ -255,6 +283,16 @@ class Pose(Detect):
     def forward(self, x):
         """Perform forward pass through YOLO model and return predictions."""
         bs = x[0].shape[0]  # batch size
-        kpt = torch.cat([self.cv4[i](x[i]).view(bs, self.nk, -1) for i in range(self.nl)], -1)  # (bs, 17*3, h*w)
+        if torch.onnx.is_in_onnx_export():
+            kpt = [self.cv4[i](x[i]).permute(0, 2, 3, 1).unsqueeze(1) for i in range(self.nl)]
+        else:
+            kpt = torch.cat([self.cv4[i](x[i]).view(bs, self.nk, -1) for i in range(self.nl)], -1)  # (bs, 17*3, h*w)
         x = self.detect(self, x)
+
+        if torch.onnx.is_in_onnx_export():
+            output = []
+            for i in range(self.nl):
+                output.append((torch.cat([x[i], kpt[i]], dim=-1)))
+            return output

Create a python file written as follows to export onnx model.

export.py

from ultralytics import YOLO
model = YOLO("./runs/pose/train/weights/best.pt")
results = model.export(format="onnx")

Enter rknn-toolkit2/examples/onnx/yolov5 and modify test.py as follows.

test.py

# Create RKNN object
rknn = RKNN(verbose=True)
 
# pre-process config
print('--> Config model')
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3588')
print('done')
 
# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model='./yolov8n_pose.onnx')
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')
 
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
if ret != 0:
    print('Build model failed!')
    exit(ret)
print('done')
 
# Export RKNN model
print('--> Export rknn model')
ret = rknn.export_rknn('./yolov8n_pose.rknn')
if ret != 0:
    print('Export rknn model failed!')
    exit(ret)
print('done')

Run test.py to generate rknn model.

$ python3 test.py

Clone the source code from our khadas/edge2-npu.

$ git clone https://github.com/khadas/edge2-npu

$ sudo apt update
$ sudo apt install cmake libopencv-dev

Put yolov8n_pose.rknn in edge2-npu/C++/yolov8n_pose/data/model.

# Compile
$ bash build.sh
 
# Run
$ cd install/yolov8n_pose
$ ./yolov8n_pose data/model/yolov8n_pose.rknn data/img/bus.jpg

Put yolov8n_pose.rknn in edge2-npu/C++/yolov8n_pose_cap/data/model.

# Compile
$ bash build.sh
 
# Run USB camera
$ cd install/yolov8n_pose_cap
$ ./yolov8n_pose_cap data/model/yolov8n_pose.rknn usb 60
 
# Run MIPI camera
$ cd install/yolov8n_pose_cap
$ ./yolov8n_pose_cap data/model/yolov8n_pose.rknn mipi 42