[이론 정리] [torchvision] 자동으로 이미지 데이터셋 생성해주는 함수

2023. 3. 16. 13:37AI/이론 정리

torchvision.datasets.ImageFolder

root/
    ├── class1/
    │   ├── img1.jpg
    │   ├── img2.jpg
    │   ├── ...
    ├── class2/
    │   ├── img3.jpg
    │   ├── img4.jpg
    │   ├── ...
    ├── ...
import torchvision.datasets as datasets

# 이미지 데이터셋 경로 설정
data_path = 'root/'

# ImageFolder 클래스를 사용하여 데이터셋 생성
dataset = datasets.ImageFolder(data_path, transform=transforms.ToTensor())

(source: https://openai.com)

 

class별로 폴더를 나누고 하위 디렉토리에 img 파일들을 넣어주고 torchvision.datasets.ImageFolder 함수를 통해 dataset을 생성할 수 있다.

이 dataset[0]은 image 정보를 담고 있고 dataset[1]은 label 정보를 담고 있다.