본문 바로가기
MM(MyMemo)

[Pytorch] 이미지 데이터셋을 torchvision을 사용해서 증가 시키기

by Glory_Choi 2023. 6. 24.
반응형
import cv2
import os
import torch
import torchvision
from torchvision.utils import save_image
file_path = './image/train/'
filenames = os.listdir('./image/train/')

for filename in filenames:
    img = cv2.imread(os.path.join(file_path, filename), cv2.IMREAD_COLOR)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    transform = torchvision.transforms.ToTensor()
    img = transform(img)
    transform = torchvision.transforms.CenterCrop(500)
    img = transform(img)

    save_image(img, os.path.join('./image/trans/',filename))

 

반응형