Pytorch 사용자 정의 Transforms.Compose 만들기

- 우리가 원하는 대로 데이터를 핸들링 하려면 Pytorch 에서 제공해주는 Compose 기능을 사용할 수 있다.

 

- 간단히 Cifar 10 의 데이터를 가지고 우리가 원하는 대로 rescale, randomCrop , 그리고 이를 합친 Compose([rescale,randomcrop]) 을 구현해보겠다.

 

- 우선 원하는 모듈을 로딩시켜주고,

import torch
import torch.nn as nn
import numpy as np
import torch.optim as optim

import torchvision
import torchvision.transforms as transforms

import matplotlib.pyplot as plt
%matplotlib inline

 

- Cuda가 잘 돌아가는지 확인 ~

device = 'cuda' if torch.cuda.is_available() else 'cpu'

torch.manual_seed(777)
if device =='cuda':
    torch.cuda.manual_seed_all(777)

 

 

- 기본적으로 돌아가는 torchvision.transfroms 의 모듈을 이용하면, transform 을 구현할 수 있다. 하지만 여기서는 더 나아가서 custom 하는 방법을 알아보자.

transform1 = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./cifar10', train=True,
                                        download=True, transform=transform1)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=512,
                                          shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR10(root='./cifar10', train=False,
                                       download=True, transform=transform1)

testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=0)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

 

- trainset 을 불러와 transpose 시켜준다. 이는 imshow 가 ( : , : ,3) 의 파일을 읽는데, 우리의 데이터 셋은 (3, : , :) 로 되어있기 때문이다. 실행시켜주게 되면,

sample = trainset[1][0].numpy()
sample = numpy.transpose(sample,(1,2,0))
plt.imshow(sample)
type(sample)

 

- 잘 뜨는거 확인했으면, ~ 이제 custom class를 작성하자. class를 작성해야 나중에 쉽게 관리가능하기 때문. Pytorch는 객체지향 너무나 좋아한다...

class Rescale(object):
    """주어진 사이즈로 샘플크기를 조정합니다.

    Args:
        output_size(tuple or int) : 원하는 사이즈 값
            tuple인 경우 해당 tuple(output_size)이 결과물(output)의 크기가 되고,
            int라면 비율을 유지하면서, 길이가 작은 쪽이 output_size가 됩니다.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, int)
        self.output_size = output_size

    def __call__(self, sample):
        image= sample
        
		# 파일의 높이 width 확인
        h, w = image.shape[:2]
        if isinstance(self.output_size, int):
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size
		
        new_h, new_w = int(new_h), int(new_w)
		# skimage.transform 이용해 resize 시켜준다.
        img = transform.resize(image, (new_h, new_w))
        return img


class RandomCrop(object):
    """샘플데이터를 무작위로 자릅니다.

    Args:
        output_size (tuple or int): 로 원하는 크기만큼 자르자.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, int)
        if isinstance(output_size, int):
        	# output_size 를 정사각형 튜플 형태로 맞춰주고,
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    def __call__(self, sample):
        image= sample

        h, w = image.shape[:2]
        #각 각 new_h,new_w 로 설정
        new_h, new_w = self.output_size
		
        # randint 를 이용해 랜덤하게 자른다.
        top = np.random.randint(0, h - new_h)
        left = np.random.randint(0, w - new_w)

        image = image[top: top + new_h,
                      left: left + new_w]


        return  image

 

 

- 이렇게 작성완료 했으면, 요롷게 객체화 시켜주고, transforms 로 묶어준다. transforms로 묶게되면, 원하는 작업들을 모조리 수행 여기서는 Rescale(24) 와 RandomCrop(18) 을 수행한 값을 반환.

- 그리고, Rescale 과 RandomCrop 같이 Compose 로 묶어주려면, 같은 객체를 반환, 입력 받아야 된다. 즉 Rescale 의 반환값을 RandomCrop 이 받을 수 있어야 한다는 것이다.( 안그러면 , 프로그램이 천재도 아니고, 어찌 자동수행한단 말이오..)

scale = Rescale(24)
crop = RandomCrop(18)
composed = transforms.Compose([Rescale(24),
                               RandomCrop(18)])

 

- 그래서 이제 값 실행해주면,

- 이쁘게 나온거 확인했으면, (위에가 Rescale, 아래가 composed 값)

- 최종으로

fig = plt.figure()
k = [scale, crop, composed]

ax = plt.subplot(1, 4, 1)
ax.set_title("sample")
plt.imshow(sample)

for i, tsfrm in enumerate(k):
    transformed_sample = tsfrm(sample)

    ax = plt.subplot(1, 4, i + 2)
    plt.tight_layout()
    ax.set_title(type(tsfrm).__name__)
    plt.imshow(transformed_sample)

plt.show()

-실행시켜주면, 이쁘게 나온 우리의 값을 확인 가능하겠다.

- 왼쪽부터 원본, 리스케일, 랜덤으로 자른것, 리스케일 후 랜덤자르기 이다.