Keypoint detection - 사람의 위치와 자세를 추정한다

Collins Lesulie on Unsplash

 

컴퓨터가 사람의 위치와 자세를 자동으로 분석해 준다면, 무엇이 좋을까요? 상상력을 발휘해 보면, 먼저 사람들의 자세를  자동으로 분석하여 교정하는 프로그램을 떠올릴 수 있겠습니다. 또한, 보행자를 분석하는 것이 중요한 자율 주행 기술에서도 요긴하게 사용될 것 같네요.

 

이외에도 다양한 분야에 사용될 수 있을것 같은데, 그 활용도를 고려하여 이미 해당 기술을 다루는 다양한 모델이 개발되어 왔습니다. 이번 프로젝트에서는 그 중의 하나를 활용하여 사진 속 사람의 위치와 자세를 추정해 보겠습니다. 

 

torchvision의 'KeypointRCNN' 모델은 이미지에서 사람의 위치를 추정하고, 신체의 keypoint를 검출하도록 훈련된 모델입니다. 이 글에서는 KeypointRCNN 모델 활용과, matplotlib을 이용하여 원하는 figure를 그리는 것을 연습해 보겠습니다. 먼저, KeypointRCNN이 어떻게 작동하는지 알아보기 위해, source code를 뜯어보겠습니다.  

 

 

 

 

---

 

 

 

 

 Source Code

link : https://github.com/pytorch/vision/blob/main/torchvision/models/detection/keypoint_rcnn.py

class KeypointRCNN(FasterRCNN):
...

def keypointrcnn_resnet50_fpn(
    ...
):	
    backbone = resnet50(weights=weights_backbone, progress=progress, norm_layer=norm_layer)
    backbone = _resnet_fpn_extractor(backbone, trainable_backbone_layers)
    model = KeypointRCNN(backbone, num_classes, num_keypoints=num_keypoints, **kwargs)

 

KeypointRCNN 클래스는 FasterRCNN 클래스를 계승하고 있으며, keypointrcnn_resnet50_fpn모델은 resnet-50을 backbone으로 하여 KeypointRCNN 클래스를 호출하고 있습니다. 

 

즉, 해당 모델은 resnet-50을 backbone으로 하는 FasterRCNN 네트워크인 것입니다. 우리가 이용할 것은 바로 이 모델로, torchvision에서는 이 모델에 대해 이미 잘 훈련된 weight를 제공하고 있습니다. 

 

 

model output

 

모델은 boxes, labels, scores, keypoints, keypoints_scores를 output으로 산출합니다. 

 

 - boxes는 사람의 위치를 box 형태로 추정합니다. (N, 4)의 tensor 형태인데, N은 box( = 사람으로 추정되는 물체)의 개수이며, 각 box마다 (x1, y1, x2, y2)의 좌표를 추정합니다. (x1, y1)은 박스 왼쪽 하단의 좌표를, (x2, y2)는 박스 오른쪽 상단의 좌표를 의미합니다.

 

 - labels는 (N, 1)의 형태로, 각 box 안의 물체의 label을 추정하는데, 이 모델에서는 전부 인간 (label = 1)으로 추정되므로 별 의미가 없는 output입니다. 

 

 - scores는 (N, 1)의 형태로, 각 box마다 인간이라고 믿어지는 정도를 나타냅니다. inference를 몇 번 시행해 보니, 0.95 이상의 score를 가진 box는 대부분 인간의 형태를 잘 추정하고 있었습니다. 

 

 - keypoints는 (N, 17, 3)의 형태로, N개의 box마다 총 17개의 keypoints를 (x, y, visiblity)의 형태로 추정합니다. 1~5번째는 눈,코,귀의 위치를 추정하고, 6~11번째는 팔의 위치(좌우 어깨, 팔꿈치, 손)를, 12~17번째는 다리의 위치를 추정합니다. 

 

 - keypoints_scores는 (N, 17)의 형태로, N개의 box마다 17개의 keypoints의 신뢰도를 추정합니다. 

 

 

 

 

---

 

 

 

 

Let's do it!

사용할 모델이 어떻게 작동하는지 알아보았으니, 이제 해당 모델을 활용해서 inference를 진행해 보겠습니다. 먼저, 필요한 라이브러리들을 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

사람의 팔다리를 연결하는 선을 그리기 위해, matplotlib.path의 Path라는 함수를 이용하는데, 해당 함수에서는 선의 모양을 결정하기 위해 codes라는 변수를 받아 활용합니다. 각 좌표마다 하나의 code를 받는데, Path.MOVETO는 선의 시작점을, Path.LINETO는 다음 점으로의 직선 연결을 의미합니다. 

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

 

이제 Gpu를 사용하도록 device를 지정하고, 모델을 evaluation mode로 불러온 후 추정할 이미지를 불러옵니다. 

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = models.detection.keypointrcnn_resnet50_fpn(pretrained=True).to(device).eval()
path = 'drive/MyDrive/projects/person_keypoint/data/person_7.jpg'
img = Image.open(path)

불러온 이미지는 image 객체이므로, gpu 환경에서 사용 가능한 적당한 크기의 tensor 객체로 변환하는 transform을 작성합니다. 

trf = T.Compose([
    T.Resize(500),
    T.ToTensor()
])
out = model([trf(img).to(device)])[0]

out 변수에는 사람의 위치와 keypoints에 대한 정보가 담겨 있습니다. 이제 바탕 이미지를 그리고, out에 담긴 정보를 바탕으로 그림을 그리겠습니다. 

figure, axes = plt.subplots()
img_trf = trf(img).permute(1,2,0).numpy()
axes.imshow(img_trf)

threshold = 0.95

for box, score, keypoints in zip(out['boxes'], out['scores'], out['keypoints']):
  score = score.detach().cpu().numpy()
  if score < threshold:
    continue

  #person detection box
  box = box.detach().cpu().numpy()
  rect = patches.Rectangle((box[0],box[1]),box[2]-box[0],box[3]-box[1],linewidth=2, edgecolor='white', facecolor='none')
  axes.add_patch(rect)

  #keypoints detection
  keypoints = keypoints.detach().cpu().numpy()[:,:2]
  for i in range(5):
    face_point = patches.Circle(keypoints[i],radius = 3, facecolor = 'yellow')
    axes.add_patch(face_point)

  leftarm_path = Path(keypoints[5:10:2],codes)
  line = patches.PathPatch(leftarm_path, linewidth=2, facecolor='none', edgecolor='red')
  axes.add_patch(line)

  rightarm_path = Path(keypoints[6:11:2],codes)
  line = patches.PathPatch(rightarm_path, linewidth=2, facecolor='none', edgecolor='red')
  axes.add_patch(line)

  leftleg_path = Path(keypoints[11:16:2],codes)
  line = patches.PathPatch(leftleg_path, linewidth=2, facecolor='none', edgecolor='red')
  axes.add_patch(line)

  rightleg_path = Path(keypoints[12:17:2],codes)
  line = patches.PathPatch(rightleg_path, linewidth=2, facecolor='none', edgecolor='red')
  axes.add_patch(line)

  for i in range(5,17):
    body_point = patches.Circle(keypoints[i],radius = 5, facecolor = 'red')
    axes.add_patch(body_point)

box에는 pathes.Rectangle 객체, 각 keypoint에는 patches.Circle 객체, 팔다리를 잇는 선에는 patches.PathPatch 객체를 각각 만들어 표시했습니다. 이 작업은 score가 0.95 이상인 box들에 대해서만 진행했습니다. 

 

 

 

 

---

 

 

 

 

마치며

이 프로젝트에서는 기본적인 pytorch 활용과 (모델 불러오기, 이미지 변환) matplotlib를 이용하여 이미지에 원하는 도형을 그리는 것을 학습했습니다. 또한, 모델의 source code를 뜯어보니, KeypointRCNN은 상위 클래스인 FasterRCNN을 기반으로 작동하고 있다는 것도 알게 되었습니다. 다음에는 FasterRCNN에 대해서도 공부해 봐야겠네요. 

 

또한, 이 프로젝트는 다음의 포스팅을 아주 많이 참고하였습니다. 공부에 도움을 주셔서 정말 감사합니다!

http://www.gisdeveloper.co.kr/?p=8235

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