사람의 눈, 코, 귀 그리고 팔과 다리를 검출하는 Person Keypoints Detection

이미지에 대한 Detection의 한 종류로 Person Keypoints Detection이 있습니다. 이 Detection은 사람의 눈, 코, 귀 그리고 팔과 다리를 검출합니다. 아래처럼요.

머신러닝 라이브리 중에 하나인, PyTorch에서는 Person Keypoints Detection을 위한 모델에 대해서 이미 잘 학습된 데이터를 제공합니다. 이 글은 미리 학습된 데이터를 활용하여 이미지에서 사람의 눈, 코, 귀와 팔 그리고 다리를 검출하는 코드를 살펴봅니다.

아울러 이 글에서는 matplotlib에서 그래프나 이미지를 단순히 표시하는 것에서 원하는 도형을 원하는 위치에 표시하는 API도 파악할 수 있습니다.

먼저 필요한 패키지를 import 합니다.

import numpy as np
from PIL import Image

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

import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches

미리학습된 데이터를 활용하여 Person Keypoints Detection을 위한 모델을 가져오는데, 머신러닝은 대규모의 행렬연산에 최적화된 GPU에서 수행하는 것이 제맛이므로 GPU 연산을 통해 신경망 모델을 로드합니다.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # CPU to GPU
model = models.detection.keypointrcnn_resnet50_fpn(pretrained=True).to(device).eval()

이미지를 불러오고 이미지의 크기를 줄입니다. GPU의 메모리가 충분하다면 이미지를 원본 크기 그대로 사용해도 되지만, 이미지를 가로로 800px로 줄여줍니다. 메모리를 덜 사용하는 것도 있고, 속도도 크게 향상되겠죠.

IMAGE_SIZE = 800
 
img = Image.open('data/running_1.jpg')
img = img.resize((IMAGE_SIZE, int(img.height * IMAGE_SIZE / img.width)))

검출을 수행합니다. 이를 위해 먼저 입력 이미지를 텐서로 변환해야합니다.

trf = T.Compose([
    T.ToTensor()
])

input_img = trf(img).to(device) # CPU to GPU

out = model([input_img])[0]

out 변수에 사람에 대한 눈, 코, 귀 그리고 팔과 다리에 대한 정보가 담겨 있습니다. 이 변수를 이용해 이미지 상에 해당 정보를 시각화 합니다.

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

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

THRESHOLD = 0.9 # 해당 정보의 정확도가 90% 이상인 것만 사용

for box, score, keypoints in zip(out['boxes'], out['scores'], out['keypoints']):
    score = score.detach().cpu().numpy() # GPU to CPU

    if score < THRESHOLD:
        continue

    box = box.detach().cpu().numpy() # GPU to CPU
    keypoints = keypoints.detach().cpu().numpy()[:, :2] # GPU to CPU

    # 사람에 대한 영역을 그리기
    rect = patches.Rectangle((box[0], box[1]), box[2]-box[0], box[3]-box[1], linewidth=2, edgecolor='white', facecolor='none')
    ax.add_patch(rect)

    # 왼쪽 팔에 대한 선 그리기
    path = Path(keypoints[5:10:2], codes)
    line = patches.PathPatch(path, linewidth=2, facecolor='none', edgecolor='red')
    ax.add_patch(line)
    
    # 오른쪽 팔에 대한 선 그리기
    path = Path(keypoints[6:11:2], codes)
    line = patches.PathPatch(path, linewidth=2, facecolor='none', edgecolor='red')
    ax.add_patch(line)
    
    # 왼쪽 다리에 대한 선 그리기
    path = Path(keypoints[11:16:2], codes)
    line = patches.PathPatch(path, linewidth=2, facecolor='none', edgecolor='red')
    ax.add_patch(line)
    
    # 오른쪽 다리에 대한 선 그리기
    path = Path(keypoints[12:17:2], codes)
    line = patches.PathPatch(path, linewidth=2, facecolor='none', edgecolor='red')
    ax.add_patch(line)

    # 눈,코,귀는 노란색으로 그리고 팔다리의 시작점과 끝점은 빨간색으로 그리기
    for i, k in enumerate(keypoints):
        if i < 5:
            RADIUS = 5
            FACE_COLOR = 'yellow'
        else:
            RADIUS = 10
            FACE_COLOR = 'red'

        circle = patches.Circle((k[0], k[1]), radius=RADIUS, facecolor=FACE_COLOR)
        ax.add_patch(circle)    

plt.show()    

어제 휴일이라 일찍 잠들었더니, 새벽 3시에 일어났습니다. 글을 작성하던 중에 새벽 4시 조금 넘어서 마켓컬리 배송이 왔네요. 시간을 초월한 마켓컬리의 배송.. 고생이 많으십니당..