PyTorch와 Numpy에서의 이미지 데이터 포맷 차이

일반적으로 이미지 데이터는 너비, 높이, 채널 순서로 표현되는 RGB 포맷을 사용한다.

예를 들어, 32x32 크기의 컬러 이미지의 경우 (32,32,3) 으로 표현되며, 각 차원은 각각 (행, 열, 채널) 을 나타낸다. 

이미지를 (행,열,채) 형식으로 읽고 처리하는 라이브러리는 다음과 같다.

• Numpy
• OpenCV
• Matplotlib (이미지 처리는 아니지만 데이터 시각화 할 때)

 

이를 확인해보기 위해 numpy 로 랜덤한 정수를 생성하여 3채널 (RGB) 이미지를 생성하는 코드를 작성해 보았다. 

생성한 이미지 데이터의 shape을 확인하면 (행,열,채널) 순서로 출력되는 것을 확인할 수 있다. 

import numpy as np
import matplotlib.pyplot as plt

# 가상의 이미지 데이터 생성 (높이: 4, 너비: 4, 채널: 3)
a = 5
b = 4
c = 3
image_data = np.random.randint(0, 256, (a, b, c), dtype=np.uint8)
print(image_data.shape)

# Matplotlib를 사용하여 이미지 표시
plt.imshow(image_data)
plt.axis('off')  # 축 제거
plt.show()

(5,4,3)

코드 출력

 

반면 PyTorch의 텐서는 이미지 데이터를 읽을 때 (채널, 행, 열)로 읽는다. 

Matplotlib의 imshow() 함수는 행, 열, 채널 순서로 이미지를 표시하므로 PyTorch 텐서를 그대로 imshow() 함수에 전달하면 이미지 데이터를 인식하지 못해 다음과 같은 오류가 뜬다.

"TypeError: Invalid shape (1, 3, 100, 150) for image data"

import torch
import matplotlib.pyplot as plt

x_batch = torch.rand(1, 3, 100, 150)

plt.imshow(x_batch)
plt.imshow(x_batch[0].permute(1, 2, 0))
plt.show()

이를 해결하기 위해, 'permute(1,2,0)'로 차원을 조정해주어 PyTorch 텐서를 Matplotlib에 적합한 순서로 변경해줘야 한다.

import torch
import matplotlib.pyplot as plt

x_batch = torch.rand(1, 3, 4, 4)

plt.imshow(x_batch[0].permute(1, 2, 0))
plt.show()