[torch] Custom dataset 만들기 (labelling 과정 포함)

질문 1: 보통 이미지당 label은 어떤 형식으로 저장돼?

• 이미지 데이터의 라벨은 다양한 방식으로 저장될 수 있으며, 사용하는 데이터셋과 상황에 따라 달라집니다. 일반적인 방법은 다음과 같습니다:

1. CSV 파일 또는 JSON 파일:

• 이미지 파일 이름과 해당 이미지의 라벨을 매핑한 파일
• 예를 들어, CSV 파일에서 각 행은 이미지 파일 이름과 그에 해당하는 라벨을 포함할 수 있습니다.

2. 폴더 구조:

• 이미지를 라벨별로 서로 다른 폴더에 저장
• 폴더 이름이 라벨을 나타내며, 이미지 파일은 해당하는 라벨의 폴더 안에 위치

3. 파일 이름:

• 이미지 파일 이름 자체에 라벨 정보를 포함시키는 방법
• 예를 들어, cat_001.jpg는 'cat' 클래스의 이미지임을 나타냅니다. 이 방법은 파일 이름 파싱을 필요로 합니다.

질문 2: 이미지를 가진 폴더를 라벨링 하기 위해, 라벨 1에 해당하는 파일명을 input으로 넣어주면, 너가 답변해준 질문1에 해당하는 형식으로 labelled data를 저장해주는 코드 작성해줘!

• 이 예시에서는 이미지 파일 이름과 라벨을 CSV 파일 형식으로 저장하는 방법을 선택
• 파일 이름에 따라 라벨 0 또는 1을 할당하며, 라벨 1에 해당하는 파일명을 input으로 받아 처리합니다.

import csv
import os


def save_labels_to_csv(image_folder: str, label_1_filenames: list,
                       csv_path: str):
    """
    이미지 폴더와 라벨 1에 해당하는 파일명 리스트를 받아, CSV 파일에 이미지와 라벨을 저장합니다.
    """
    with open(csv_path, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(['filename', 'label'])  # CSV 헤더

        for filename in os.listdir(image_folder):
            label = 1 if filename in label_1_filenames else 0
            writer.writerow([filename, label])


# 예시 사용법
image_folder = 'path/to/your/images'
label_1_filenames = ['image1.jpg', 'image2.jpg']  # 라벨 1에 해당하는 이미지 파일명
csv_path = 'path/to/your/labels.csv'
save_labels_to_csv(image_folder, label_1_filenames, csv_path)

질문 3: 그렇게 저장된 labelled data를 불러와서, CustomDataset 만드는 코드 작성해줘.

• CSV 파일에서 라벨링된 데이터를 불러와 CustomDataset 클래스를 생성하는 방법을 설명
• 먼저, CSV 파일을 읽어 이미지 파일 경로와 라벨을 로드한 다음, 이 정보를 사용해 CustomDataset 인스턴스를 생성

import pandas as pd
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, Resize, ToTensor
from PIL import Image


class CustomDataset(Dataset):

    def __init__(self, image_folder: str, csv_file: str, transform=None):
        self.df = pd.read_csv(csv_file)
        self.image_folder = image_folder
        self.transform = transform

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.image_folder, self.df.iloc[idx, 0])
        image = Image.open(img_name).convert('RGB')
        label = int(self.df.iloc[idx, 1])
        if self.transform:
            image = self.transform(image)
        return image, label


# 데이터셋 및 데이터 로더 준비
image_folder = 'path/to/your/images'
csv_file = 'path/to/your/labels.csv'
transform = Compose([
    Resize((224, 224)),  # CLIP 모델에 맞는 이미지 크기 조정
    ToTensor()  # 이미지를 텐서로 변환
])

# CustomDataset 인스턴스 생성
dataset = CustomDataset(image_folder=image_folder,
                        csv_file=csv_file,
                        transform=transform)

# DataLoader 인스턴스 생성
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

# 데이터 로더를 사용한 반복 예시
for images, labels in dataloader:
    # 여기서 이미지와 라벨을 모델에 공급하는 처리를 수행할 수 있습니다.
    pass

• 이 코드는 먼저 CustomDataset 클래스를 정의하여 초기화 함수에서 이미지 폴더 경로, 라벨이 저장된 CSV 파일 경로, 그리고 전처리 변환을 인수로 받습니다.
• len 메서드는 데이터셋의 크기를 반환하고, getitem 메서드는 인덱스에 해당하는 이미지와 라벨을 로드한 후, 전처리 변환을 적용하여 반환합니다.
• DataLoader는 이 데이터셋을 배치 크기로 나누고, 필요에 따라 데이터를 섞어서 모델 학습에 사용할 준비를 합니다.
• 이 구조를 사용하면 데이터셋에 있는 이미지들을 효율적으로 로드하고 전처리하여 PyTorch 모델에 공급할 수 있습니다.