Torchvision 을 사용한 사람의 키포인트를 찾아내는 딥러닝 모델 예제

Torchvision 을 사용하여 사람의 키포인트를 찾아내는 딥러닝 모델 예제 입니다.

Torchvision 의 keypointrcnn_resnet50_fpn 모델을 참조하세요.

 

 

Torchvision 을 사용한 사람의 키포인트 를 찾아내는 딥러닝 모델 예제

import torch
import torchvision
from torchvision import models
import torchvision.transforms as T

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches

IMG_SIZE = 480
THRESHOLD = 0.95

---

# Pytorch keypointrcnn Model Download
model = models.detection.keypointrcnn_resnet50_fpn(pretrained=True).eval()

# Load Image Resizing
img = Image.open('Images/group-of-people-2297627.jpg')
img = img.resize((IMG_SIZE, int(img.height * IMG_SIZE / img.width)))

# Image to Tensor Conversion
trf = T.Compose([
       T.ToTensor()
])

input_img = trf(img)
print(input_img.shape)

 

# Inference
out = model([input_img])[0]

▣ Load 한 이미지를 보다 작게 resize 처리 한 결과와 keypointrcnn_resnet50_fpn 모델 결과 이미지 정보를 보여줍니다.

---

codes = [
    Path.MOVETO,
    Path.LINETO,
    Path.LINETO
]

fig, ax = plt.subplots(1, figsize=(16, 16))
ax.imshow(img)

▣ 모델 결과 이미지정보중 box 와 score ,  keypoint만을 사용한 결과를 보여줍니다.

▣ box 는 사람을 인식한 위치 포인트를 나타내고, score 는 사람을 인식한 확률값입니다.

▣ keypoint 는 인식되 사람의 핵심 point를 나타냅니다.

for box, score, keypoints in zip(out['boxes'], out['scores'], out['keypoints']):
    # Numpy Array 로 변환한다.
    score = score.detach().numpy()

    ▣ score가 THRESHOLD (0.95) 이하인경우 제외 시킵니다.
    if score < THRESHOLD:
        continue

    box = box.detach().numpy()
    keypoints = keypoints.detach().numpy()[:, :2]

    rect = patches.Rectangle((box[0], box[1]), box[2]-box[0], box[3]-box[1], linewidth=2,
             edgecolor='b', facecolor='none')
    ax.add_patch(rect)

    # 17 keypoints
    for k in keypoints:
        circle = patches.Circle((k[0], k[1]), radius=2, facecolor='r')
        ax.add_patch(circle)

    # draw path
    # left arm
    path = Path(keypoints[5:10:2], codes)
    line = patches.PathPatch(path, linewidth=2, facecolor='none', edgecolor='r')
    ax.add_patch(line)

    # right arm
    path = Path(keypoints[6:11:2], codes)
    line = patches.PathPatch(path, linewidth=2, facecolor='none', edgecolor='r')
    ax.add_patch(line)

    # left leg
    path = Path(keypoints[11:16:2], codes)
    line = patches.PathPatch(path, linewidth=2, facecolor='none', edgecolor='r')
    ax.add_patch(line)

    # right leg
    path = Path(keypoints[12:17:2], codes)
    line = patches.PathPatch(path, linewidth=2, facecolor='none', edgecolor='r')
    ax.add_patch(line)

plt.show()

 

실행 결과

 

이미지 출처 : https://www.pexels.com/search/group%20of%20people/

Torchvision 을 사용한 사람의 키포인트를 찾아낸 모델 실행 결과

파란색 박스의 사람을 인식하는 표식과 빨강 포인트와 라인으로 사람의 키 포인트 (얼굴 포인트와 신체포인트)를 인식하는 모델 결과를 확인 할 수 있습니다.