あかすくぱるふぇ

同人サークル「あかすくぱるふぇ」のブログです。

pytorch(torchvision)のdatasetsは有名なデータセットを簡単に扱うためのパッケージです。
その中にはCOCOデータセットも含まれているのですが、セグメンテーション情報については(元々のデータセットの時点で)符号化されていて、そのままでは利用することができません。
そこで本記事では、pytorchでCOCOデータセットのセグメンテーション情報をデコードして取得する方法を解説します。

まずは、普通にCOCOデータセットからデータをLoadする方法。
import torch
from torch.utils import data
from torchvision import transforms, datasets

if __name__ == '__main__':

# DataLoaderを生成
transform = transforms.Compose([
transforms.ToTensor()
])
train_set = datasets.CocoDetection(root='../data/coco/train2017',
annFile='../data/coco/instances_train2017.json',
transform=transform)
train_loader = torch.utils.data.DataLoader(train_set)

# DataLoad
for data, target in train_loader:
print('hoge')
データセットはあらかじめCOCOデータセットのwebサイトからダウンロードしておきます。
ImagesとAnnotationsをセットでダウンロードしてください。
http://cocodataset.org/#download

上記のコードで画像dataとアノテーションデータtargetをLoadできます。
ただし、前述したように、targetの'segmentation'キーの値は符号化されていて、そのままでは利用できません。

続いて、セグメンテーション情報のデコードです。
Loadしてからデコードしてもよいのですが、tensor型に変換されてLoadされてくるので、処理が煩雑になってしまいます。
そこで、CocoDetectionを継承したクラスCocoSegmentationを作って、その内部でセグメンテーション情報をデコードしてLoadするようにします。
以下がCocoSegmentationクラスのコードです。
from torchvision.datasets import coco
from pycocotools.mask import frPyObjects, decode

class CocoSegmentation(coco.CocoDetection):
def __getitem__(self, index):
img, target = super().__getitem__(index)
for category in target:
seg_rle = category['segmentation']
category['segmentation'] = decode(frPyObjects(seg_rle, img.shape[1], img.shape[2]))
return img, target
'segmentation'キーの値をデコードして書き換えているだけです。
デコード処理については以下のページが参考になります。
https://github.com/cocodataset/cocoapi/issues/4

あとは、最初に示したコードでCocoDetectionの代わりにCocoSegmentationクラスを利用すれば、デコードされたセグメンテーション情報を利用することができます。

以上です。



・追記
上のコードだとリサイズやバッチ処理が入ると動かないことが判明しました。
最終的なコードは以下のようになりました。
import numpy as np
from torchvision import transforms
from torchvision.datasets import coco
from pycocotools.mask import frPyObjects, decode
from PIL import Image
import os

class CocoSegmentation(coco.CocoDetection):
def __getitem__(self, index):

# データ入力
coco = self.coco
img_id = self.ids[index]
ann_ids = coco.getAnnIds(imgIds=img_id)
target = coco.loadAnns(ann_ids)
path = coco.loadImgs(img_id)[0]['file_name']
img = Image.open(os.path.join(self.root, path)).convert('RGB')

# セグメンテーション情報のデコード
for category in target:
seg_rle = category['segmentation']
tmp = decode(frPyObjects(seg_rle, img.size[1], img.size[0]))
if tmp.ndim == 3:
tmp = np.sum(tmp, axis=2, dtype=np.uint8)
category['segmentation'] = tmp

# data_transform
if self.transform is not None:
img = self.transform(img)

# target_transform
for category in target:
pilImg = Image.fromarray(category['segmentation'])
tmp = pilImg.resize((img.shape[2], img.shape[1]), resample=Image.NEAREST)
target_transform = transforms.Compose([
transforms.ToTensor()
])
category['segmentation'] = target_transform(tmp)

return img, target

本記事では、踊ってみた動画から人の3Dポーズを抽出する方法を解説します。
結果として以下のような3Dポーズを抽出することができます。


抽出処理には、以下の2つの技術を利用しています。

・OpenPose(RGB画像から人の2Dポーズを推定)
https://arxiv.org/abs/1611.08050
https://github.com/CMU-Perceptual-Computing-Lab/openpose

・3d-pose-baseline(2Dポーズから3Dポーズへの変換)
https://arxiv.org/abs/1705.03098
https://github.com/una-dinosauria/3d-pose-baseline


作業手順は以下の通りです。
1. 元動画を連続静止画に変換
2. 各静止画からOpenPoseで関節の二次元位置を抽出
3. 関節の二次元位置を時間方向に平滑化
4. 関節の二次元位置を3d-pose-baselineの入力形式に変換
5. 3d-pose-baselineで関節の三次元位置を推定

2〜5について、詳細に説明していきます。



2. 各静止画からOpenPoseで関節の二次元位置を抽出

まず、上記githubからOpenPoseを持ってきてQuickStartが動くところまで進めてください。
その後、examples/tutorial_pose/1_extract_from_image.cppに手を加えることで、
関節の二次元位置をファイル出力できるようにします。

関節の二次元位置はposeKeypointsに格納されています。
const auto poseKeypoints = poseExtractorCaffe.getPoseKeypoints();
poseKeypointsは[人][関節位置とスコア]の二次元配列です。
動画に一人しか映っていないのであれば、一次元目のサイズは1になっていてほしいのですが、
実際には人でないものを人と認識してしまい、一次元目のサイズが1より大きくなっていることが多いです。

そこで、まずは二次元目に格納されているスコアを使って、主被写体を特定します。
二次元目が(関節1のx座標, 関節1のy座標, 関節1のスコア, 関節2のx座標, 関節2のy座標, 関節2のスコア, ...)となっているので、全関節のスコアを確認し、最も合計スコアが高い人を主被写体として特定します。
愚直に書くとこんな感じ。
const auto numberKeypoints = poseKeypoints.getSize(1);
double max_score = 0;
int max_score_person = 0;
for (auto person = 0 ; person < poseKeypoints.getSize(0) ; ++person) {
int index = person * numberKeypoints * 3;
double score = 0;
for (auto point = 0; point < numberKeypoints; ++point) {
score += poseKeypoints[index + 2];
index += 3;
}
if (score >= max_score) {
max_score = score;
max_score_person = person;
}
}
あとは、主被写体の関節の二次元位置をファイル出力すれば完了です。



3. 関節の二次元位置を時間方向に平滑化

OpenPoseの出力した二次元位置は突発的な大きな誤差を含んでいるので、時間方向に平滑化します。
上の動画を作った際は、各関節について3フレームの座標値の中央値によって平滑化しました。



4. 関節の二次元位置を3d-pose-baselineの入力形式に変換

OpenPoseと3d-pose-baselineでは関節の二次元位置の形式が異なっているので変換します。

OpenPoseの出力形式は以下のページに載ってます。
https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/output.md

3d-pose-baselineの入力形式はdata_utils.pyのH36M_NAMESを参照してください。

以上を踏まえて、OpenPose形式から3d-pose-baseline形式への変換は以下のコードで実現できます。
lがOpenPose形式、enc_inが3d-pose-baseline形式です(3d-pose-baseline/predict_3dpose.py参照)。
order = [15, 12, 25, 26, 27, 17, 18, 19, 1, 2, 3, 6, 7, 8]
for
i in range(len(order)):
for j in range(2):
enc_in[0][order[i] * 2 + j] = l[i * 2 + j]

for j in range(2):
# Hip
enc_in[0][0 * 2 + j] = (enc_in[0][1 * 2 + j] + enc_in[0][6 * 2 + j]) / 2
# Neck/Nose
enc_in[0][14 * 2 + j] = (enc_in[0][15 * 2 + j] + enc_in[0][12 * 2 + j]) / 2
# Thorax
enc_in[0][13 * 2 + j] = 2 * enc_in[0][12 * 2 + j] - enc_in[0][14 * 2 + j]



5. 3d-pose-baselineで関節の三次元位置を推定

3d-pose-baseline/predict_3dpose.pyで関節の三次元位置を推定します。
ただ、predict_3dpose.pyには以下のような問題があります。
・元画像のx, y方向への全身の動き(オフセット)が無視される
・3の平滑化処理で取りきれなかった大きな誤差で出力ポーズが暴れる
・今回の目的に関係ない余計な処理が多い

以上の問題を解決するために、ぐちゃぐちゃといじった結果のコードを以下に載せておきます。
汚すぎて自分でも何書いてるのかよくわかりません\(^o^)/

#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import tensorflow as tf
import data_utils
import viz
import cameras
from predict_3dpose import create_model

FLAGS = tf.app.flags.FLAGS
order = [15, 12, 25, 26, 27, 17, 18, 19, 1, 2, 3, 6, 7, 8]

def main(_):

enc_in = np.zeros((1, 64))
enc_in[0] = [0 for i in range(64)]

actions = data_utils.define_actions(FLAGS.action)

SUBJECT_IDS = [1, 5, 6, 7, 8, 9, 11]
rcams = cameras.load_cameras(FLAGS.cameras_path, SUBJECT_IDS)
train_set_2d, test_set_2d, data_mean_2d, data_std_2d, dim_to_ignore_2d, dim_to_use_2d = data_utils.read_2d_predictions(
actions, FLAGS.data_dir)
train_set_3d, test_set_3d, data_mean_3d, data_std_3d, dim_to_ignore_3d, dim_to_use_3d, train_root_positions, test_root_positions = data_utils.read_3d_data(
actions, FLAGS.data_dir, FLAGS.camera_frame, rcams, FLAGS.predict_14)

device_count = {"GPU": 1}
with tf.Session(config=tf.ConfigProto(
device_count=device_count,
allow_soft_placement=True)) as sess:
batch_size = 128
model = create_model(sess, actions, batch_size)

for fileNo in range(1100):

print(fileNo)

csvName = 'csv_median/median{0:06d}.csv'.format(fileNo+1)
f = open(csvName, 'r')

for line in f:
line = line.rstrip()
l = line.split(",")
for i in range(len(order)):
for j in range(2):
enc_in[0][order[i] * 2 + j] = l[i * 2 + j]

for j in range(2):
# Hip
enc_in[0][0 * 2 + j] = (enc_in[0][1 * 2 + j] + enc_in[0][6 * 2 + j]) / 2
# Neck/Nose
enc_in[0][14 * 2 + j] = (enc_in[0][15 * 2 + j] + enc_in[0][12 * 2 + j]) / 2
# Thorax
enc_in[0][13 * 2 + j] = 2 * enc_in[0][12 * 2 + j] - enc_in[0][14 * 2 + j]

f.close()

spine_x = enc_in[0][24]
spine_y = enc_in[0][25]

enc_in = enc_in[:, dim_to_use_2d]
mu = data_mean_2d[dim_to_use_2d]
stddev = data_std_2d[dim_to_use_2d]
enc_in = np.divide((enc_in - mu), stddev)

dp = 1.0
dec_out = np.zeros((1, 48))
dec_out[0] = [0 for i in range(48)]
_, _, poses3d = model.step(sess, enc_in, dec_out, dp, isTraining=False)

enc_in = data_utils.unNormalizeData(enc_in, data_mean_2d, data_std_2d, dim_to_ignore_2d)
poses3d = data_utils.unNormalizeData(poses3d, data_mean_3d, data_std_3d, dim_to_ignore_3d)

gs1 = gridspec.GridSpec(1, 1)
gs1.update(wspace=-0.00, hspace=0.05) # set the spacing between axes.
plt.axis('off')

subplot_idx, exidx = 1, 1

max = 0
min = 10000
for i in range(poses3d.shape[0]):
for j in range(32):
tmp = poses3d[i][j * 3 + 2]
poses3d[i][j * 3 + 2] = poses3d[i][j * 3 + 1]
poses3d[i][j * 3 + 1] = tmp
if poses3d[i][j * 3 + 2] > max:
max = poses3d[i][j * 3 + 2]
if poses3d[i][j * 3 + 2] < min:
min = poses3d[i][j * 3 + 2]

for i in range(poses3d.shape[0]):
for j in range(32):
poses3d[i][j * 3 + 2] = max - poses3d[i][j * 3 + 2] + min
poses3d[i][j * 3] += (spine_x - 630)
poses3d[i][j * 3 + 2] += (500 - spine_y)

# Plot 3d predictions
ax3 = plt.subplot(gs1[subplot_idx - 1], projection='3d')

if np.min(poses3d) < -1000:
poses3d = before_pose

p3d = poses3d
viz.show3Dpose(p3d, ax3, lcolor="#9b59b6", rcolor="#2ecc71")

pngName = 'png/test{0:06d}.png'.format(fileNo+1)
plt.savefig(pngName)

before_pose = poses3d

if __name__ == "__main__":
tf.app.run()
以上です。

↑このページのトップヘ