# Visualize 32x32 image to 224x224
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt

original_img = train_ds[0]['img']

# 이미지 크기를 224x224로 변환하는 변환기를 정의한다.
#CIFAR-10은 32*32의 이미지데이터셋인데, 크기를 resize해줌
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

# 이미지를 변환 -> Transform을 이용해서 사이즈 키워주기
transformed_img = transform(original_img)

# 텐서를 이미지로 변환
transformed_img = transforms.ToPILImage()(transformed_img)

# 이미지를 시각화
plt.imshow(transformed_img)
plt.axis('off')  # 축을 표시하지 않기

from torchvision.transforms import (CenterCrop,
                                    Compose,
                                    Normalize,
                                    RandomHorizontalFlip,
                                    RandomResizedCrop,
                                    Resize,
                                    ColorJitter,
                                    GaussianBlur,
                                    RandomGrayscale,
                                    ToTensor)

normalize = Normalize(mean=image_mean, std=image_std) #이미지 normalize
_train_transforms = Compose(
        [
# train_dataset에 다양한 Augmentation 적용하기
            Resize(size),
            RandomHorizontalFlip(),
            RandomResizedCrop(size),
            ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
            GaussianBlur(kernel_size=3),
           #RandomGrayscale(p=0.2),
            ToTensor()         
    )

_val_transforms = Compose(
        [
            Resize(size),
            CenterCrop(size),
            ToTensor(),
            normalize,
        ]
    )

# 모든 train과 val 데이터에 대해서 위에서 정의한 transform을 적용하기
def train_transforms(examples):
    examples['pixel_values'] = [_train_transforms(image.convert("RGB")) for image in examples['img']]
    return examples

def val_transforms(examples):
    examples['pixel_values'] = [_val_transforms(image.convert("RGB")) for image in examples['img']]
    return examples

viz_data = train_ds[3]
original_image = viz_data['img']
augmented_image = viz_data['pixel_values']

transform = transforms.Resize((224, 224))
original_image = transform(original_image)
augmented_image = transform(augmented_image)

fig, axes = plt.subplots(1,2, figsize=(10,5))

axes[0].imshow(original_image) #첫번째 공간에 original_image 그리기
axes[0].set_title("Original Image") #제목설정
axes[0].axis('off') #불필요한 축 숨기기

axes[1].imshow(augmented_image.permute(1,2,0)) #두번째 공간에 augmented_image 그리기
axes[1].set_title("Augmented Image") #제목설정
axes[1].axis('off') #불필요한 축 숨기기

plt.tight_layout() #subplot간 간격 자동조절

plt.show()