사람의 눈, 코, 귀 그리고 팔과 다리를 검출하는 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시 조금 넘어서 마켓컬리 배송이 왔네요. 시간을 초월한 마켓컬리의 배송.. 고생이 많으십니당..

“사람의 눈, 코, 귀 그리고 팔과 다리를 검출하는 Person Keypoints Detection”에 대한 7개의 댓글

  1. 정보 감사합니다.
    궁금한점은 out 결과물이 나오는데 10 이상 걸립니다.
    혹시 모델 중 빠른 처리가 되는 모델 추천 가능한가요?

  2. 질문 하나 남김니다.
    인터넷이 되는 PC에서 아래 라인을 실행하면
    에러 없이 잘 실행됩니다.
    그리고 같은 PC에서 랜선 또는 환경 설정에서 인터넷 미사용하여
    실행 해도 잘 실행 됩니다.

    그런데 아예 인터넷이 안되는 PC(망분리) 에서
    아래 라인을 실행하면
    urllib.error.URLError: 라고 에러 발생합니다.
    제 생각에는 아래 코드는 인터넷과 무관한 것 같은데..
    고수님 의견 부탁합니다.
    감사합니다.

    — 아래 —
    model = models.detection.keypointrcnn_resnet50_fpn(pretrained=True).eval()

  3. 감사합니다.

    result_out = model([tensor_img])[0]
    위 모델 생성에서 약 6초 정도 걸립니다.
    이미지 사이즈는 582 * 462 입니다.
    PC 사양이 좋지 않고 CPU 만 사용합니다.

    이 정도 나오는게 정상인가요?

    같은 이미지로 mediapipe로 할때는 1초 이하로 진행됩니다.

답글 남기기

이메일 주소는 공개되지 않습니다. 필수 필드는 *로 표시됩니다