drivingDataset.py

import os
import pandas as pd
from PIL import Image
import torch
from torch.utils.data import Dataset

# DrivingDataset 클래스는
# Unity 에서 만든 (PNG 이미지 + labels.csv)를
# PyTorch 가 바로 학습 가능한 형태 (Tensor 이미지 + Float 라벨)로 변환하는 클래스

# class DrivingDataset(Dataset)
# PyTorch torch.utils.data.Dataset 을 상속한 커스텀 데이터셋 클래스를 정의
#   - Dataset 을 상속하면
#     __len__() 과 __getitem__() 을 반드시 구현해야 함
#   - 이 클래스는 
#     Unity에서 생성한 이미지 + csv 라벨을 PyTorch 학습 데이터로 변환하는 역학
class DrivingDataset(Dataset):
    # Dataset 객체 초기화
    def __init__(self, root_dir, transform=None):
        # 이미지와 csv 가 저장된 폴더 경로 저장
        # "DrivingDataset" 로 경로 지정 필요
        self.root_dir = root_dir

        # DrivingDataset 폴더 안의 label.csv 를 pandas 로 읽음
        # csv의 구조는 
        # frame,steering,throttle,break
        # 1,0.0000,0.0000,0.0000
        # 2,0.0000,0.0000,0.0000
        # 3,0.0000,0.0000,0.0000
        # 4,0.0000,0.0000,0.0000
        # self.labels 는 pandas DataFrame 형태
        self.labels = pd.read_csv(os.path.join(root_dir, "labels.csv"))

        # 이미지를 PyTorch 학습 데이터로 변환하는 transform 저장
        # 일반적으로
        # transform.Compose([
        #   transform.Resize((224, 224)),
        #   trandform.ToTensor(),
        # ])
        self.transform = transform

    # Dataset 전체 샘플 수를 리턴
    def __len__(self):
        # label.csv 에 총 10,000 줄이면 -> length=10,000
        # DataLoader 가 batch를 만들 때 사용
        return len(self.labels)
    
    # Dataset의 핵심
    # index 번째 샘플(이미지 + 라벨)을 반환하는 함수
    # PyTorch는 내부적으로 아래처럼 호출
    # img, label = dataset[5]
    def __getitem__(self, index):
        # labels.csv 에서 index 번째 행을 가져옴
        # ex> index=100 이면 csv 100번째 줄
        row = self.labels.iloc[index]

        # row 에서 frame 번호를 읽음
        # 유니티에서 저장한 png 파일 이름의 번호와 동일
        # ex> frame_000100.png
        frama_idx = int(row["frame"])

        # 실제 이미지 파일 경로 작성
        # {frame_idx:06d} 는 zero-padding(6자리)
        # ex> frame_100 -> frame-000100.png
        img_path = os.path.join(self.root_dir, f"frame_{frama_idx:06d}.png")

        # 이미지 불러오기
        # PIL(Pillow) 로 PNG 파일을 읽음
        # convert("RGB")로 3채널 고정 (PyTorch 입력은 3채널 필요)
        image = Image.open(img_path).convert("RGB")

        # transform 이 설정되어 있으면 적용
        # Resize, Tensor 변환, Normalize 등이 여기에서 실행
        # ex> ToTensor() -> (H, W, C) 이미지를 FloatTensor(C, H, W) 로 변환 
        # Normalize -> 모델 학습에 적합한 값으로 조정
        if self.transform:
            image = self.transform(image)

        # 라벨 구성
        # 회귀(regression) 학습을 위해 float32 텐서로 만듬
        # 라벨 구조 : [ steering, throttle, brake ]
        # ex> [-0.7, 1.0, 0.0]
        # 즉, 이 모델은 이미지 -> 이 3개 값 회귀를 목표로 함
        # steering, throttle, brake -> 회귀용 라벨
        label = torch.tensor([
            row["steering"],
            row["throttle"],
            row["break"]
        ], dtype=torch.float32)

        # 최종 출력
        # image -> PyTorch Tensor(C, H, W)
        # label -> Tensor(3)
        # DataLoader 는 이를 batch 단위로 묶어서 모델의 입력으로 제공
        return image, label