rtmpose

前后处理代码

import cv2
import numpy as np
import onnxruntime

map = [(15, 13), (13, 11), (16, 14), (14, 12), (11, 12), (5, 11),
       (6, 12), (5, 6), (5, 7), (6, 8), (7, 9), (8, 10), (1, 2),
       (0, 1), (0, 2), (1, 3), (2, 4), (3, 5), (4, 6)]


def preprocess(img, input_size):
    h, w = img.shape[:2]
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    scale = min(input_size[0] / w, input_size[1] / h)
    x, y = (input_size[0] - w * scale) / 2, (input_size[1] - h * scale) / 2
    img = cv2.warpAffine(img, np.array([[scale, 0, x], [0, scale, y]]), input_size)
    img = (img - [123.675, 116.28, 103.53]) / [58.395, 57.12, 57.375]
    img = img.transpose(2, 0, 1)
    return img, scale, (x, y)


def postprocess(outputs, scale, top_left):
    simcc_x, simcc_y = outputs[0][0], outputs[1][0]
    x_locs, y_locs = np.argmax(simcc_x, axis=1), np.argmax(simcc_y, axis=1)
    x_locs, y_locs = (x_locs / 2 - top_left[0]) / scale, (y_locs / 2 - top_left[1]) / scale
    max_val_x, max_val_y = np.amax(simcc_x, axis=1), np.amax(simcc_y, axis=1)
    mask = max_val_x > max_val_y
    max_val_x[mask] = max_val_y[mask]
    point = np.transpose([x_locs, y_locs, max_val_x])
    return point


def draw_img(img, keypoint, p):
    for i, j in map:
        if keypoint[i][2] > p and keypoint[j][2] > p:
            p1 = np.int32(keypoint[i][:2])
            p2 = np.int32(keypoint[j][:2])
            cv2.line(img, p1, p2, (255, 255, 0), 1, cv2.LINE_AA)
    for i in keypoint:
        if i[2] > p:
            cv2.circle(img, np.int32(i[:2]), 2, (0, 255, 0), 3, cv2.LINE_AA)



if __name__ == '__main__':
    model_path = '../model/rtm-pose-t.onnx'
    image = cv2.imread('../img/aa.jpg')
    model = onnxruntime.InferenceSession(model_path)
    input_name = model.get_inputs()[0].name
    img, scale, top_left = preprocess(image, (192, 256))
    outputs = model.run(None, {input_name: [img]})
    keypoint = postprocess(outputs, scale, top_left)
    draw_img(image, keypoint, 0.4)
    cv2.imshow("a", image)
    cv2.waitKey(0)