matplotlib 做 MPII 人体 skeleton

2019-07-14  本文已影响0人  谢小帅
005808361.jpg 005808361.jpg

annotation

a = {
    "joints_vis": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],  # 16 joints
    "joints": [
        [804.0, 711.0], [816.0, 510.0], [908.0, 438.0], [1040.0, 454.0],
        [906.0, 528.0], [883.0, 707.0], [974.0, 446.0], [985.0, 253.0],
        [982.7591, 235.9694], [962.2409, 80.0306], [869.0, 214.0], [798.0, 340.0],
        [902.0, 253.0], [1067.0, 253.0], [1167.0, 353.0], [1142.0, 478.0]
    ],
    "image": "005808361.jpg",
    "scale": 4.718488,
    "center": [966.0, 340.0]
}

vis_skeleton.py

import torch
import numpy as np
import skimage.io
import cv2
import pylab as plt
import math
import time
import os

# 16 joints
kp_names = [
    'r_ankle', 'r_knee', 'r_hip', 'l_hip', 'l_knee', 'l_ankle',  # legs, [0-5]
    'pelvis', 'spine', 'neck', 'head',  # middle body, [6-9]
    'r_wrist', 'r_elbow', 'r_shoulder', 'l_shoulder', 'l_elbow', 'l_wrist',  # arms, [10-15]
]
num_joints = len(kp_names)

# 15 limbs
limbSeq = [
    [0, 1], [1, 2], [2, 6], [6, 3], [3, 4], [4, 5],  # legs
    [6, 7], [7, 8], [8, 9],  # middle body
    [10, 11], [11, 12], [12, 8], [8, 13], [13, 14], [14, 15]  # arms
]
num_limbs = len(limbSeq)

colors = [
    [0, 255, 0], [0, 255, 85], [0, 255, 170], [255, 255, 0], [170, 255, 0], [85, 255, 0],
    [255, 170, 0], [0, 255, 0], [255, 85, 0],
    [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], [170, 0, 255],
    [255, 0, 85]
]


def plt_skeleton(img_path, joints):
    plt.figure(figsize=(12, 7.2))
    # plt.axis('off')
    image = plt.imread(img_path)
    plt.imshow(image)
    plt.title(os.path.basename(img_path) + '\n')

    # plt joints
    for k in range(num_joints):  # 16
        plt.plot(joints[k][0], joints[k][1], marker='o', color=[c / 255 for c in colors[k]])
        plt.text(joints[k][0], joints[k][1], '{}'.format(k),
                 size=14, color='w')
    # plt bones
    for i in range(num_limbs):
        k_A, k_B = limbSeq[i]
        plt.plot([joints[k_A][0], joints[k_B][0]],
                 [joints[k_A][1], joints[k_B][1]],
                 color=[c / 255 for c in colors[i]], linestyle='-', linewidth=2)
    # save
    plt.savefig(img_path.replace('imgs', 'res'),
                bbox_inches="tight",  # 保存的结果和显示的结果基本一致
                # pad_inches=0.0)
                )
    plt.show()


ann = {
    "joints_vis": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],  # 16 joints
    "joints": [
        [804.0, 711.0], [816.0, 510.0], [908.0, 438.0], [1040.0, 454.0],
        [906.0, 528.0], [883.0, 707.0], [974.0, 446.0], [985.0, 253.0],
        [982.7591, 235.9694], [962.2409, 80.0306], [869.0, 214.0], [798.0, 340.0],
        [902.0, 253.0], [1067.0, 253.0], [1167.0, 353.0], [1142.0, 478.0]
    ],
    "image": "005808361.jpg",
    "scale": 4.718488,
    "center": [966.0, 340.0]
}

plt_skeleton(img_path=os.path.join('./imgs', ann['image']),
             joints=ann['joints'])
上一篇下一篇

猜你喜欢

热点阅读