PyTorch参考書 OpenPoseのサンプルコードの追加フラグメント

「つくりながら学ぶ! PyTorchによる発展ディープラーニング」の「第4章 姿勢推定（OpenPose）」の「Stageモジュール」の各Stageの出力を可視化するフラグメントを作成しました。

参照書籍

つくりながら学ぶ! PyTorchによる発展ディープラーニング
小川雄太郎

https://book.mynavi.jp/ec/products/detail/id=104855
https://github.com/YutaroOgawa/pytorch_advanced
https://qiita.com/sugulu/items/07253d12b1fc72e16aba

目的

Stageが複数ある理由の調査（lossを複数回算出する理由。特徴量を複数回入力する理由。）。
⇒1段目～6段目で、推定結果がどの程度異なるか、目視で確認する。

各Stageの出力可視化フラグメント

「4-7_OpenPose_inference.ipynb」を以下の通り、修正・追記ください。

セル修正前

predicted_outputs, _ = net(x)

セル修正後

predicted_outputs, intermediate_outputs = net(x)

セル追加

def convert_to_numpy_and_resize(val, size, size2):
    # 画像をテンソルからNumPyに変化し、サイズを戻します
    val = val.detach().numpy().transpose(1, 2, 0)
    val = cv2.resize(val, size, interpolation=cv2.INTER_CUBIC)
    val = cv2.resize(val, size2, interpolation=cv2.INTER_CUBIC)
    return val

セル追加

size2 = (oriImg.shape[1], oriImg.shape[0])
pafs_list = []
heatmaps_list = []
for i in range(6):
    pafs_list.append(convert_to_numpy_and_resize(intermediate_outputs[i*2][0], size, size2))
    heatmaps_list.append(convert_to_numpy_and_resize(intermediate_outputs[i*2+1][0], size, size2))

セル追加

def display_overlay_img(featImg, oriImg, anoNo, i):
    featImg = featImg[:, :, anoNo]
    featImg = Image.fromarray(np.uint8(cm.jet(featImg)*255))
    featImg = np.asarray(featImg.convert('RGB'))

    # 合成して表示
    blend_img = cv2.addWeighted(oriImg, 0.5, featImg, 0.5, 0)

    plt.subplot(1,3,i+1)
    plt.imshow(blend_img)
    # plt.show()

def display_heatmap_and_paf(heatmaps, pafs, oriImg, anoNos):
    plt.figure(figsize=(6.4*3, 4.3*1))
    display_overlay_img(heatmaps, oriImg, anoNos[0], 0)
    display_overlay_img(heatmaps, oriImg, anoNos[1], 1)
    display_overlay_img(pafs, oriImg, anoNos[2], 2)
    plt.show()

セル追加

print("左肘heatmap、左手首heatmap、PAF")
for i in range(6):
    print("Stage" + str(i+1))
    display_heatmap_and_paf(heatmaps_list[i], pafs_list[i], oriImg, (6, 7, 24))

セル追加

print("右膝heatmap、右足首heatmap、PAF")
for i in range(6):
    print("Stage" + str(i+1))
    display_heatmap_and_paf(heatmaps_list[i], pafs_list[i], oriImg, (9, 10, 5))

出力結果（1例目：左肘heatmap、左手首heatmap、PAF）

出力結果（2例目：右膝heatmap、右足首heatmap、PAF）