rknn调用人体骨架点api

import cv2
import time
from rockx import RockX
import numpy as np
FACE_DETECT_MIN_SCORE = 0.6
#存放了模型和数据
skeleton = [[12,13],[9,10],[14, 13],[10, 11],[0,1],[1,2],[1,6],[1,3],[6,7],[3,4],[4,5],[7,8],[2,12],[2,9]]#一共16个点
ROCKX_DATA_DIR = "/home/coolpi/RockX_SDK_V1.6.0_20220913/sdk/rockx-data-rk3588"
handle_face = RockX(RockX.ROCKX_MODULE_FACE_DETECTION, rockx_data_dir=ROCKX_DATA_DIR)
handle_face_landmarks = RockX(RockX.ROCKX_MODULE_FACE_LANDMARK_5, rockx_data_dir=ROCKX_DATA_DIR)
handle_body_landmarks = RockX(RockX.ROCKX_MODULE_POSE_BODY, rockx_data_dir=ROCKX_DATA_DIR)
#调用RockX类的rockx_person_detect方法
process_face = getattr(handle_face, "rockx_face_detect")
# process_map = {
#         RockX.ROCKX_MODULE_FACE_DETECTION : ("face", "rockx_face_detect", FACE_DETECT_MIN_SCORE),
#         RockX.ROCKX_MODULE_PERSON_DETECTION : ("body", "rockx_person_detect", BODY_DETECT_MIN_SCORE),
#         RockX.ROCKX_MODULE_HEAD_DETECTION : ("head", "rockx_head_detect", HEAD_DETECT_MIN_SCORE),
#     }
process_face_landmark = getattr(handle_face_landmarks, "rockx_face_landmark") #rockx_pose_body
process_body_landmark = getattr(handle_body_landmarks, "rockx_pose_body")
facepath = "datatest/run/"
dataset = 'datatest/stand1.mp4'
cap = cv2.VideoCapture(dataset)
flag, frame = cap.read()
j = 0
det_results = [] 
while flag:
    # 读取每一帧，flag表示是否读取成功，frame为图片内容。
    flag, frame = cap.read()
    # frame = cv2.pyrDown(frame)
    if frame is None:
        # print("No video captured")
        continue
    else:
        width = frame.shape[1]
        height = frame.shape[0]
        ################
        ret, results = process_face(
            frame, width, height,RockX.ROCKX_PIXEL_FORMAT_BGR888)       
        # results = self._filter_results(results)
        ################
        # ret, results = process(frame, width, height,RockX.ROCKX_PIXEL_FORMAT_BGR888)
        for result in results:
            if result.score >FACE_DETECT_MIN_SCORE:
                t1 = time.time()
                b = [result.box.top, result.box.right, result.box.bottom, result.box.left]
                ret, face_landmark = process_face_landmark(frame, width, height, RockX.ROCKX_PIXEL_FORMAT_BGR888, result.box)
                face_landmarks = face_landmark[4]
                for landmark in face_landmarks:
                        cv2.circle(frame, (landmark.x,landmark.y), 3, (255, 0, 0), 3)
                ret, keypoints_array = process_body_landmark(frame, width, height, RockX.ROCKX_PIXEL_FORMAT_BGR888, original_ratio=1)
                # cv2.line(img, start_point, end_point, (0, 255, 0), 2)
                # cv2.rectangle(frame, (b[3],b[0]), (b[1],b[2]), (0, 0, 0), -1)
                for keypoint_array in keypoints_array:
                    body_landmarks = keypoint_array[2]
                    i = 0
                    for landmark in body_landmarks:
                        cv2.circle(frame, (landmark.x,landmark.y), 3, (255, 0, 0), 3)
                        cv2.putText(frame,str(i),(landmark.x,landmark.y),cv2.FONT_HERSHEY_SIMPLEX,1,(0,255,255),3)
                        i = i + 1
                    for sk in skeleton:
                        x0 = body_landmarks[sk[0]].x
                        y0 = body_landmarks[sk[0]].y
                        x1 = body_landmarks[sk[1]].x
                        y1 = body_landmarks[sk[1]].y
                        cv2.line(frame, (x0, y0), (x1, y1),(0, 255, 0), 100)
                t2 = time.time()
                print("时间消耗：",t2-t1,"s")
                cv2.imwrite(facepath+str(j)+'.jpg', frame)
                j=j+1
                print(j)