[파이토치] Custom Dataset 만들기

https://github.com/stonegyoung/Pytorch_Study/blob/master/catdog_custom_dataset%20%EB%A7%8C%EB%93%A4%EA%B8%B0.ipynb

Pytorch_Study/catdog_custom_dataset 만들기.ipynb at master · stonegyoung/Pytorch_Study

---

custom dataset을 만들기 위해서는 세 가지 방법이 있습니다.

1. 클래스 별로 폴더 구성하고 관리

train/test 분리와 라벨링이 따로 필요하다

(데이터셋의 크기가 작거나 간단한 구조일 때 적합)

2. train, test 폴더 안에 클래스 별 폴더 구성하고 관리

train/test 분리가 되어 있다

(데이터셋 크기가 크고, train/test 비율이 사전 설정된 경우에 적합)

3. 통합 파일 기반 관리

train, test 파일에 이미지 파일 경로, 라벨을 저장한다.

(대규모 데이터셋일 경우 적합)

 

---

1. 클래스 별로 폴더 구성하고 관리

1. torch.utils.data의 Dataset 상속받는 클래스 생성
2. init, len, getitem 메소드를 만든다
- init: 경로 설정, train과 test 분리, 정답 데이터 리스트 만들기
- len: 전체 데이터셋 개수 리턴
- getitem: 특정 번호의 이미지와 정답 리턴

 

필요 라이브러리 임포트

import torch
from torch.utils.data import Dataset
import glob
from PIL import Image

 

 

현재 디렉토리 상태

 

 

glob 

 

custom dataset 클래스 생성

class CatDog(Dataset):
  def __init__(self, path, train=True, transform=None): # path: 이미지가 있는 경로 train: train인지 test인지 transform: 사용할 transform
    self.path = path # ./catdog
    self.transform = transform

    # glob으로 이미지 경로 불러오기
    self.cat_list = glob.glob(path+'/cat/*.*')
    self.dog_list = glob.glob(path+'/dog/*.*')
    
    # train, test 분리 + 정답 데이터 리스트 만들기
    # train 70 test 30
    # 이미지 70퍼센트 가져오기
    cat_idx = int(len(self.cat_list) * 0.7)
    dog_idx = int(len(self.dog_list) * 0.7)
    if train == True:
      self.data_list = self.cat_list[:cat_idx] + self.dog_list[:dog_idx]
      # 정답 리스트 만들기(cat 0, dog 1)
      self.label_list = [0] * cat_idx + [1] * dog_idx
    else:
      self.data_list = self.cat_list[cat_idx:] + self.dog_list[dog_idx:]
      self.label_list = [0] * len(self.cat_list[cat_idx:]) + [1] * len(self.cat_list[dog_idx:]) # 0은 cat, 1은 dog
      
  def __len__(self):
    return len(self.data_list)

  def __getitem__(self, idx):
    img_path = self.data_list[idx] # data_list에는 경로가 적혀 있음
    label = self.label_list[idx] # 정답
    img = Image.open(img_path) # 경로를 통해 이미지 불러오기

    if self.transform is not None: # transform 있으면
      img = self.transform(img) # 전처리한 이미지

    return img, label

 

train/test 분리가 필요하고, 라벨링도 따로 해줘야 한다.

 

전처리

data_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((224,224))
])

 

train_data, test_data 생성

path = './catdog'
train_data = CatDog(path=path, train=True, transform=data_transform)
test_data = CatDog(path=path, train=False, transform=data_transform)

 

---

2. train, test 폴더 안에 클래스 별 폴더 구성하고 관리

1. train과 test 폴더에 각각 카테고리 별 폴더를 넣어준다

현재 디렉토리 상태 -> train, test 디렉토리로 나누기

필요 라이브러리 임포트

import os
import shutil

 

이미지 복사 함수

# 이미지 복사
def dataset_split(folder, train_cnt=40):
  directory_list = ['./animal/train', './animal/test']
  img_path_list = os.listdir(f'./catdog/{folder}')
  for directory in directory_list:
    if not os.path.isdir(directory+'/'+folder):
      os.makedirs(directory+'/'+folder)
  for cnt in range(len(img_path_list)):
    filename = img_path_list[cnt]
    if cnt < train_cnt: # train 폴더에 복사
      print(f'{filename} train 폴더에 복사')
      shutil.copy(f'./catdog/{folder}/{filename}', f'./animal/train/{folder}/{filename}')  
    else: # test 폴더에 복사
      print(f'{filename} test 폴더에 복사')
      shutil.copy(f'./catdog/{folder}/{filename}', f'./animal/test/{folder}/{filename}')

 

 

torchvision의 datsets의 ImageFolder를 이용한다

 

필요 라이브러리 임포트

import torch
from torchvision.datasets import ImageFolder
from torchvision import transforms

 

전처리

data_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((224,224))
])

 

train_data, test_data 생성(ImageFolder 사용)

train_data = ImageFolder(root='./animal/train', transform=data_transform)
test_data = ImageFolder(root='./animal/test', transform=data_transform)

 

속성 살펴보기

.classes

.class_to_idx

.imgs

 

---

3. 통합 파일 기반 관리

csv 파일을 이용하여 custom dataset을 만들겠습니다.

csv는 train.csv, test.csv로 나누고,

그 안에 있어야 할 것

1. 이미지 파일 경로 file_path

2. 정답 클래스인 label

 

일단 2번에서 만든 구조를 사용하겠습니다.

train.csv / test.csv

csv 파일을 읽어서 이미지를 출력해보겠습니다.

 

필요 라이브러리 임포트

import pandas as pd
from PIL import Image

 

pandas를 이용해 csv 파일 불러오기

train_df = pd.read_csv('train.csv')
test_df = pd.read_csv('test.csv')

train_df.head()

 

file_path의 경로를 PIL.Image로 읽어서 출력해보겠습니다

Image.open(train_df.loc[0, 'file_path'])

이미지가 잘 나온다면, 준비 끝

 

 

1번처럼 custom dataset을 만들어주겠습니다.

 

필요 라이브러리 임포트

import torch
from torchvision import transforms # 전처리
from torch.utils.data import Dataset, DataLoader # 데이터셋, 데이터로더
import matplotlib.pyplot as plt # 시각화
import numpy as np

 

전처리

data_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((224,224))
])

 

custom dataset 클래스 생성

 

class CatDog(Dataset):
  def __init__(self, df, transform=None): # df: file_path와 label이 담긴 데이터프레임, transform: 사용할 transform
    self.data_list = df['file_path'].values # 넘파이 배열로
    self.label_list = df['label'].values # 넘파이 배열로
    self.transform = transform

  def __len__(self):
    return len(self.data_list)

  def __getitem__(self, idx):
    img_path = self.data_list[idx] # data_list에는 경로가 적혀 있음
    label = self.label_list[idx] # 정답
    img = Image.open(img_path) # 경로를 통해 이미지 불러오기

    if self.transform is not None: # transform 있으면
      img = self.transform(img) # 전처리한 이미지

    return img, label

근데 path와 train 여부 없이 데이터프레임만 받으면 됩니다.

 

train_data, test_data 생성

train_data = CatDog(train_df, data_transform)
test_data = CatDog(train_df, data_transform)