아이공의 AI 공부 도전기

Pytorch에서 이미지 여는 방법

 

     

개요

단순히 어떤 폴더에 있는 이미지를 열고자 한다면 PIL을 사용하여 show 하는 방식이 가장 보편적인 방법일 것입니다.

import PIL

img = PIL.Image.open('./tree.jpg')
img.show()

해당 결과로 위와 같은 예쁜 사진을 열어볼 수 있을 것입니다. 물론 matplotlib를 통해서도 얻을 수 있겠지만 위 방법이 제일 간단하겠지요. 그러나 오늘 우리가 알아볼 것은 단순히 이미지를 여는 것이 아니라 torch기반 tensor에서 이미지를 여는 방법에 대해 알아볼 것입니다.

 

왜냐하면 torch.Tensor는 tensor로는 이미지를 바로 확인할 수 없이 한 단계를 더 거쳐야 하기 때문입니다. 그렇다면 이제 그 방법에 대해 알아봅시다.

 

Pytorch에서 이미지 여는 방법

1. PIL과 연결하여 이미지 열기

PIL을 이용하여 이미지를 열고자 하는 방법을 알려드리기 전에 우선 PIL을 torch.tensor로 변환하는 방법에 대해서 알아봅시다. 왜냐하면 해당 방법을 알게 되면 해당 방법의 역방향으로 PIL을 얻어낼 수 있기 때문입니다.

 

PIL을 torch.Tensor로 바꾸는 방법은 torchvision.transforms.ToTensor()를 통해서 진행할 수 있습니다. 해석하면 To Tensor '텐서에게로'라는 뜻이 되고 이는 즉 PIL Image를 Tensor로 바꿔주는 것이라는 것을 알 수 있습니다.

이렇게 바꾼 Tensor의 shape를 알아볼까요?

import PIL
import torchvision.transforms as transforms

img = PIL.Image.open('./tree.jpg')
#img.show()

tf = transforms.ToTensor()
img_t = tf(img)

# torch.Size([3, 1280, 1920])
print(img_t.size())

보이는 바와 같이 torch.Tensor로 잘 변환되었고 3,1280,1920이라는 channel, height, width 또한 확인할 수 있습니다.

 

자! 여기서 우리는 transform에 어떤 함수들이 존재하고 그 함수들을 통해 PIL과 tensor를 자유롭게 변환할 수 있다는 사실을 눈치챌 수 있습니다.

 

그렇다면 Tensor에서 PIL로 변환하려면 어떻게 해야할까요

바로 torchvision.transforms.ToPILImage()를 사용하면 됩니다.

tf = transforms.ToPILImage()
img_t = tf(img_t)

#<PIL.Image.Image image mode=RGB size=1920x1280 at 0x1A7B973F8C8>
print(img_t)

img_t.show()

 

TIP
 
 

위에 이어진 코드입니다. 해당 코드만 따로 실행시키지 마세요. PIL.Image의 경우 단순히 출력할 경우 위 주석과 같이만 출력되니 show를 이용하여 확인하세요.

 

자! 이렇게 코드를 돌려봄으로써 Tensor를 PIL 이미지로 PIL 이미지를 Tensor로 서로 변환할 수 있음을 확인할 수 있겠죠?

 

2. matplotlib.pyplot으로 열기

matplotlib를 통해 이미지를 연다면 아마 numpy와 관련이 제일 많을 것입니다. 그리고 torch에서는 numpy에서 제공하는 거의 대부분의 함수를 제공합니다. 때문에 numpy에서 이미지를 출력하고자 할 때 사용하는 permute를 사용하는 방법을 사용할 것입니다.

import torch
import PIL
import matplotlib.pyplot as plt
import torchvision.transforms as transforms

img = PIL.Image.open('./tree.jpg')
tf = transforms.ToTensor()
img_t = tf(img)

#torch.Size([3, 1280, 1920])
print(img_t.size())

### method 2 ###
img_t = img_t.permute(1,2,0)

#torch.Size([1280, 1920, 3])
print(img_t.size())

plt.imshow(img_t)
plt.show()

3. torchvision.utils.make_grid 사용하기

이 방법은 사실 batch가 아닌 단일 이미지에서는 권장하지는 않습니다. 다만 NN 학습에 있어 여러 개의 사진을 출력하고자 할 때 많이 사용하는 방법이기에 알아두는 편이 좋습니다.

 

사용방법은 4D 또는 list의 tensor를 넣고 해당 이미지들에 대한 행의 개수, padding, normalize를 설정하는 방법으로 결과를 얻어냅니다. 저는 단일 이미지이기 때문에 단순히 tensor만 넣는 방법으로 보여드리고자 합니다.

 

tensor_4d = img_t.unsqueeze(0)

# torch.Size([1, 3, 1280, 1920])
print(tensor_4d.size())

grid = torchvision.utils.make_grid(tensor_4d)

# torch.Size([3, 1280, 1920])
print(grid.size())

plt.imshow(grid.permute(1,2,0))
plt.show()

 

사실 2번째 방법과 매우 유사해 보입니다. 그러나 이 함수의 가장 큰 장점은 다수의 batch에서 다수의 이미지를 간단하게 표현할 수 있다는 것입니다. 바로 이 점이 장점이지만 지금 저희가 사용하는 단일 이미지에서는 큰 영향력을 발휘하지 못하는 것은 사실입니다.

 

그렇기에 상황에 맞게 본인이 좋아하는 방향대로 설정하여 사용하는 것이 좋아보입니다.

공유하기

facebook twitter kakaoTalk kakaostory naver band
loading