こんにちは　kamikawaです
今回は姿勢推定を試していきたいと思います
モデルは自分で学習させると時間がかかるので簡単に試せるtorchvisionsの学習済みモデルを使いたいと思います

実行環境はGoogle Colaboratory(Colab)です

対象読者

AIを試してみたい!
姿勢推定を試してみたい!
スポーツなどの動作解析をやってみたい!

姿勢推定(pose estimation)とは

姿勢推定とは人間の姿勢を推定するタスク(問題になります)
詳しく分けると

2D推定・・・人間の姿勢を2Dで推定する(人間の画像→棒人間みたいなイメージ)
3D推定・・・人間の姿勢を3Dで推定する(人間の画像→フィギュアみたいなイメージ)

2D推定より3D推定の方が圧倒的に難しい問題です。
今回は2D推定を試していきます

コード


import torch
import PIL
from PIL import Image
import torchvision
from torchvision import transforms
import cv2
import matplotlib.pyplot as plt

frame_raw = cv2.imread('画像のpath') 
frame = cv2.cvtColor(frame_raw,cv2.COLOR_BGR2RGB)
image = Image.fromarray(frame)

model = torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=True, progress=True, num_classes=2, num_keypoints=17, pretrained_backbone=True, trainable_backbone_layers=3)
with torch.no_grad():
  device = torch.device('cpu')
  transform = transforms.Compose([transforms.ToTensor()])
  inputs = transform(image)
  inputs = inputs.unsqueeze(0).to(device)
  model.eval()
  outputs = model(inputs)
  key_point=[]
  for i in range(len(outputs[0]['boxes'])):
    if outputs[0]['scores'][i]>=0.9:
      x0 = outputs[0]['boxes'][i][0] 
      y0 = outputs[0]['boxes'][i][1]
      x1 = outputs[0]['boxes'][i][2]
      y1 = outputs[0]['boxes'][i][3]
      if outputs[0]['scores'][i]>=0.9:
        bbox = cv2.rectangle(frame_raw,(x0,y0),(x1,y1),(0,0,300),3,4)
      x=[]
      y=[]
      for j in range(len(outputs[0]['keypoints'][0])):
        x.append(outputs[0]['keypoints'][i][j][0])
        y.append(outputs[0]['keypoints'][i][j][1])
        bbox = cv2.circle(frame_raw,(outputs[0]['keypoints'][i][j][0],outputs[0]['keypoints'][i][j][1]),2,(50,300,0),3,4)
      key_point.append((x,y))

  for a in key_point: 
    x=a[0]
    y=a[1] 
    bbox =cv2.line(frame_raw,(x[5],y[5]),(x[6],y[6]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[5],y[5]),(x[7],y[7]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[6],y[6]),(x[8],y[8]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[7],y[7]),(x[9],y[9]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[8],y[8]),(x[10],y[10]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[11],y[11]),(x[12],y[12]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[11],y[11]),(x[13],y[13]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[12],y[12]),(x[14],y[14]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[13],y[13]),(x[15],y[15]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[14],y[14]),(x[16],y[16]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[5],y[5]),(x[11],y[11]),(300,0,0),5)
    bbox =cv2.line(frame_raw,(x[6],y[6]),(x[12],y[12]),(300,0,0),5)

  plt.figure(figsize = (12,9))
  plt.imshow(cv2.cvtColor(bbox, cv2.COLOR_BGR2RGB))
  plt.axis("off")
  plt.show()

3D推定について

3D推定の論文のデモを載せておきます
単眼RGB画像から人物の3Dモデル(objectファイル)を作る研究のデモです

ひとりの人物に対しての推定
CVPR2020のPIFuHD: Multi-Level Pixel-Aligned Implicit Function for High-Resolution 3D Human Digitization

複数人の推定(障害物や人が見切れていてもOK)
ECCV2020のCenterHMR: a Bottom-up Single-shot Method
for Multi-person 3D Mesh Recovery from a Single Image

興味があればぜひ試してみてください(リンクから飛んで「Open in Colab」と書かれた青いボタン)

をクリックしてみてくださいするとColabが開きます