CHWとHWCの相互変換

やりたいこと

opencvの画像は縦(Height)×横(Width)×チャンネル(Channel)のいわゆるHWC形式になっています。
それとは異なり、pytorchの画像ではチャンネル(Channel)×縦(Height)×横(Width)のCHW形式になっています。

numpy arrayとtorch tensorにおいて、HWCとCHWの相互変換方法を示します。

Numpy Arrayでの軸変換

np.transposeで軸を変換します。

import numpy as np
H, W, C = 128, 256, 3
img = np.random.random([H, W, C])    # H, W, C
print(img.shape)

img = np.transpose(img, (2, 0, 1))        # C, H, W
print(img.shape)

実行すると次のような出力を得ます。

(128, 256, 3)  
(3, 128, 256)

ちゃんと軸変換ができました！

Pytorch Tensorでの軸変換

torch.permuteで軸変換ができます。

import torch
H, W, C = 128, 256, 3
img = torch.randn((C, H, W))         # C, H, W
print(img.size())
img = torch.permute(img, (1, 2, 0)) # H, W, C
print(img.size())

実行すると次のような出力を得ます

torch.Size([3, 128, 256])
torch.Size([128, 256, 3])

こちらもちゃんと軸変換できました！