😺

bizarre-pose-estimatorを自分のdocker環境で動かす

2024/01/22に公開

Docker

PyTorch

tech

概要

イラスト系のキャラクターのポーズ推定を行えるbizarre-pose-estimatorというモデルがあります。こちらはDockerfileが付属しているのですが、自分のdocker環境で動かしたかったので必要なモジュールをインストールして動かしました。結構大変だったのでメモとして残しておきます。動けばいい方は素直に公式のDockerfileを使ったほうが楽です。

環境構築

dockerイメージ

FROM nvidia/cuda:11.1.1-devel-ubuntu20.04

py_envのバージョン

 anaconda3-2022.05

pythonのバージョンはPython 3.9.12になります。3.10以上を使うとモジュールのインストールに失敗しました。

torchvision==0.9.1
matplotlib==3.5.0
scipy==1.7.1
scikit-learn==1.0.1
scikit-image==0.18.3
imagesize==1.3.0
flask==2.0.2
patool==1.12
easydict
pillow==8.4.0
wandb==0.12.7
pyunpack==0.2.2
opencv-contrib-python==4.5.4.60
kornia==0.6.2
pytorch-lightning==1.3.8
torchmetrics===0.6.0

追加でdetectron2をインストールします

RUN python -m pip install git+https://github.com/facebookresearch/detectron2.git

ファイルの準備

リポジトリのクローン

git clone https://github.com/ShuhongChen/bizarre-pose-estimator.git
cd bizarre-pose-estimator

モデルのダウンロード

こちらのリンクから以下の４つのzipファイルをダウンロードします。

bizarre_pose_dataset.zip
character_bg_seg_data.zip
bizarre_pose_models.zip
raw_retrieval_support.zip

zipファイルの展開と移動

zipファイルを展開して、bizarre_pose_modelsの中身をbizarre-pose-estimatorフォルダ直下に置きます。

unzip bizarre_pose_dataset.zip
unzip character_bg_seg_data.zip
unzip bizarre_pose_models.zip
unzip raw_retrieval_support.zip 

mv bizarre_pose_models/_train/character_bg_seg/* _train/character_bg_seg/
mv bizarre_pose_models/_train/character_pose_estim/* _train/character_pose_estim/
mv bizarre_pose_models/_train/danbooru_tagger/* _train/danbooru_tagger/

実行

ポーズ推定の実行

python3 -m _scripts.pose_estimator \
    ./_samples/megumin.png \
    ./_train/character_pose_estim/runs/feat_concat+data.ckpt

結果

bounding box
        top-left: (67, 12)
        size: (401, 489)

keypoints
        (212.87, 122.03) nose
        (204.46, 130.43) eye_left
        (204.46, 105.22) eye_right
        (204.46, 143.04) ear_left
        (208.67, 84.20) ear_right
        (242.29, 151.44) shoulder_left
        (254.89, 92.61) shoulder_right
        (267.50, 197.67) elbow_left
        (296.92, 96.81) elbow_right
        (254.89, 248.10) wrist_left
        (284.31, 67.39) wrist_right
        (326.33, 231.29) hip_left
        (343.14, 210.27) hip_right
        (254.89, 306.93) knee_left
        (250.69, 302.73) knee_right
        (334.74, 424.59) ankle_left
        (385.17, 336.34) ankle_right