😺
イラスト系キャラクターのポーズ推定を行ってみる
概要
イラスト系のキャラクター画像のポーズ推定をしたいのですが、事例が少ないようです。調査したところ２つのモデルを組み合わせると高い精度でポーズ推定できたのでまとめておきます。
調査
通常のモデルを探すときはPaperWithCodeで該当タスクのSOTAを調べます。ポーズ推定の場合はPose Estimationというタスクが該当します。しかしこのタスクは実写画像のみのSOTAでイラスト系のモデルを見つけることはできませんでした。実写系ポーズ推定モデルのSOTAを使えばイラスト系も扱えそうですが、実はうまくいきません。精度が上がったことでイラストは扱いにくくなっているようです。
最新のイラスト系の画像のポーズ推定について調べていたところ、こちらのページを見つけました。2023年12月に更新されているようです。
ポーズ推定の新しいモデルを探してみますが２つしか無いようです。今回は新しいほうのTransfer Learning for Pose Estimation of Illustrated Charactersを使います。こちらは以前使ったことがあり、新しいものがないか期待していたのですが存在しないようです。
今回のプラン
bizarre-pose-estimatorは以前使ったことがあるのですが、１キャラクター１画像のものしか扱えません。また背景も無いほうがよさそうです。そこで今回はCartoonSegmentationを使ってキャラクターのバウンディングボックスとマスクを取得して、キャラクターの部分だけ背景なしで取り出します。そこからbizarre-pose-estimatorでポーズ推定を行う２段構成としました。
キャラクターを一人単位で切り抜く
CartoonSegmentationを使ってキャラクターを一人単位で切り抜いて保存します。またキャラクター部分のマスクを反転して、そこに白を設定して背景を白にします。
import cv2
from PIL import Image
import numpy as np
from animeinsseg import AnimeInsSeg, AnimeInstances
from animeinsseg.anime_instances import get_color
import os
from tqdm import tqdm
# パラメータ
mask_thres = 0.3
instance_thres = 0.3
refine_kwargs = {'refine_method': 'refinenet_isnet'} # 使用しない場合はNoneを指定
# refine_kwargs = None
# 画像の準備
import glob
import random
input_dir = 'images'
output_dir = 'images_output'
os.makedirs(output_dir, exist_ok=True)
file_list = glob.glob(os.path.join(input_dir, '*'), recursive=True)
print(len(file_list))
# AnimeInsSegの準備
ckpt = r'models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt'
net = AnimeInsSeg(ckpt, mask_thr=mask_thres, refine_kwargs=refine_kwargs)
def inference(img, filename):
basename = os.path.basename(filename)
output_filename = os.path.join(output_dir, basename)
if os.path.exists(output_filename):
return
# 推論の実行
instances: AnimeInstances = net.infer(
img,
output_type='numpy',
pred_score_thr=instance_thres
)
drawed = img.copy()
im_h, im_w = img.shape[:2]
if type(instances.bboxes) != np.ndarray:
return
# 結果の確認
for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
left = int(xywh[0])
top = int(xywh[1])
right = int(xywh[2]) + left
bottom = int(xywh[3]) + top
img_clopped = np.copy(img[top : bottom, left : right])
white = np.array([255, 255, 255])
# 背景を白で塗りつぶす
mask_clopped = mask[top : bottom, left : right]
# ひっくり返す
rev_mask_clipped = np.logical_not(mask_clopped)
# 背景を白にする
img_clopped[rev_mask_clipped] = white
basename = os.path.basename(filename)
without_ext = os.path.splitext(basename)[0]
print('without_ext:', without_ext)
output_filename = os.path.join(output_dir, without_ext + '_' + str(ii) + '.png')
try:
cv2.imwrite(output_filename, img_clopped)
print('write:', output_filename)
except:
print('skip:', filename)
continue
for file in tqdm(file_list):
try:
if '.zip' in file:
continue
img = cv2.imread(file)
except:
continue
inference(img, file)
ポーズ推定
bizarre-pose-estimatorは動く状態になっているものとします。_scripts/pose_estimator.pyを複製して、指定フォルダの画像のポーズ推定を行うスクリプトを作成しました。画像ファイルをglobで取得して、ポーズ推定しているだけです。
from _util.util_v1 import * ; import _util.util_v1 as uutil
from _util.pytorch_v1 import * ; import _util.pytorch_v1 as utorch
from _util.twodee_v0 import * ; import _util.twodee_v0 as u2d
import _util.keypoints_v0 as ukey
parser = argparse.ArgumentParser()
parser.add_argument('fn_img')
parser.add_argument('fn_model')
args = parser.parse_args()
img = I(args.fn_img)
######################## SEGMENTER ########################
from _train.character_bg_seg.models.alaska import Model as CharacterBGSegmenter
model_segmenter = CharacterBGSegmenter.load_from_checkpoint(
'./_train/character_bg_seg/runs/eyeless_alaska_vulcan0000/checkpoints/'
'epoch=0096-val_f1=0.9508-val_loss=0.0483.ckpt'
)
def abbox(img, thresh=0.5, allow_empty=False):
# get bbox from alpha image, at threshold
img = I(img).np()
assert len(img) in [1,4], 'image must be mode L or RGBA'
a = img[-1] > thresh
xlim = np.any(a, axis=1).nonzero()[0]
ylim = np.any(a, axis=0).nonzero()[0]
if len(xlim)==0 and allow_empty: xlim = np.asarray([0, a.shape[0]])
if len(ylim)==0 and allow_empty: ylim = np.asarray([0, a.shape[1]])
axmin,axmax = max(int(xlim.min()-1),0), min(int(xlim.max()+1),a.shape[0])
aymin,aymax = max(int(ylim.min()-1),0), min(int(ylim.max()+1),a.shape[1])
return [(axmin,aymin), (axmax-axmin,aymax-aymin)]
def infer_segmentation(self, images, bbox_thresh=0.5, return_more=True):
anss = []
_size = self.hparams.largs.bg_seg.size
self.eval()
for img in images:
oimg = img
# img = a2bg(resize_min(img, _size).convert('RGBA'),1).convert('RGB')
img = I(img).resize_min(_size).convert('RGBA').alpha_bg(1).convert('RGB').pil()
timg = TF.to_tensor(img)[None].to(self.device)
with torch.no_grad():
out = self(timg)
ans = TF.to_pil_image(out['softmax'][0,1].float().cpu()).resize(oimg.size[::-1])
ans = {'segmentation': I(ans)}
ans['bbox'] = abbox(ans['segmentation'], thresh=bbox_thresh, allow_empty=True)
anss.append(ans)
return anss
######################## POSE ESTIMATOR ########################
if 'feat_concat' in args.fn_model:
from _train.character_pose_estim.models.passup import Model as CharacterPoseEstimator
elif 'feat_match' in args.fn_model:
from _train.character_pose_estim.models.fermat import Model as CharacterPoseEstimator
else:
assert 0, 'must use one of the provided pose estimation models'
model_pose = CharacterPoseEstimator.load_from_checkpoint(args.fn_model, strict=False)
def infer_pose(self, segmenter, images, smoothing=0.1, pad_factor=1):
self.eval()
try:
largs = self.hparams.largs.adds_keypoints
except:
largs = self.hparams.largs.danbooru_coco
_s = largs.size
_p = _s * largs.padding
anss = []
segs = infer_segmentation(segmenter, images)
for img,seg in zip(images,segs):
# segment
oimg = img
ans = {
'segmentation_output': seg,
}
bbox = seg['bbox']
cb = u2d.cropbox_sequence([
# crop to bbox, resize to square, pad sides
[bbox[0], bbox[1], bbox[1]],
resize_square_dry(bbox[1], _s),
[-_p*pad_factor/2, _s+_p*pad_factor, _s],
])
icb = u2d.cropbox_inverse(oimg.size, *cb)
img = u2d.cropbox(img, *cb)
img = img.convert('RGBA').alpha(0).convert('RGB')
ans['bbox'] = bbox
ans['cropbox'] = cb
ans['cropbox_inverse'] = icb
ans['input_image'] = img
# pose estim
timg = img.tensor()[None].to(self.device)
with torch.no_grad():
out = self(timg, smoothing=smoothing, return_more=True)
ans['out'] = out
# post-process keypoints
kps = out['keypoints'][0].cpu().numpy()
kps = u2d.cropbox_points(kps, *icb)
ans['keypoints'] = kps
anss.append(ans)
return anss
######################## MODEL FORWARD ########################
def _visualize(image=None, bbox=None, keypoints=None):
v = image
if bbox is not None:
v = v.rect(*bbox, c='r', w=2)
if keypoints is not None:
if isinstance(keypoints, dict):
keypoints = np.asarray([keypoints[k] for k in ukey.coco_keypoints])
for (a,b),c in zip(ukey.coco_parts, ukey.coco_part_colors):
v = v.line(keypoints[a], keypoints[b], w=5, c=c)
keypoints = keypoints[:len(ukey.coco_keypoints)]
for kp in keypoints:
v = v.dot(kp, s=5, c='r')
return v
import glob
from tqdm import tqdm
file_list = glob.glob('image_output/*')
for file in tqdm(file_list):
img = I(file)
ans = infer_pose(model_pose, model_segmenter, [img,])
bbox = ans[0]['bbox']
print(f'bounding box\n\ttop-left: {bbox[0]}\n\tsize: {bbox[1]}')
print()
print('keypoints')
v = img
for k,(x,y) in zip(ukey.coco_keypoints, ans[0]['keypoints']):
print((f'\t({x:.2f}, {y:.2f})'), k)
print()
_visualize(img, ans[0]['bbox'], ans[0]['keypoints']).save('./_samples/' + os.path.basename(file))
結果
こちらのCartoonSegmentationのサンプルに入っている画像を二人に分けてポーズ推定します。
元画像
分割して背景を白にした画像
ポーズ推定した結果
一部、足の位置がおかしいですが動作しているようですね。
Discussion