ACDC Dataset : PyTorch 에서 사용하기 (code)

Adverse Conditions Dataset with Correspondences

https://acdc.vision.ee.ethz.ch/

ACDC Dataset

 

ACDC 데이터셋은 다양한 시각적 조건에서의 semantic segmentation 을 위한 대규모 데이터셋입니다.

이 데이터셋에는 안개, 밤, 비, 눈과 같은 네 가지 일반적인 조건 각각에 대해 균등하게 분포된 4006 개의 이미지가 포함되어 있습니다.

사용 목적

1. semantic segmentation
2. object detection
3. panoptic segmentation

ACDC dataset 의 파일 트리

• rgb_anon_trainvaltest
  • 크기 : 15.6 GB
  • 익명화된 악조건 이미지 (Anonymized adverse-condition images)
    • fog, night, rain, snow 와 같은 4가지 주요 악조건 하에서 촬영된 이미지가 포함되어 있다.
    • 각 조건에서 촬영된 이미지는 총 4006장으로 이루어져 있다.
  • 익명화된 정상 조건 이미지 (Anonymized corresponding normal-condition images)
    • 위의 악조건 이미지와 같은 장면 / 위치에서 일반 조건 (악조건이 없는) 하에서 촬영된 이미지이다.
    • 동일하게 총 4006장으로 구성되어 있다.

좌측 : GOPR0475_frame_000247_rgb_anon.png / 우측 : GOPR0475_frame_000247_rgb_ref_anon.png

 

ACDC Dataset File Tree

 

ACDC 데이터셋은 다음과 같은 파일 트리로 구성되어 있습니다. 

 

ㄴ rgb_anon_trainvaltest
	ㄴ fog 
		ㄴ test
			ㄴ GOPR0475
            ㄴ GOPR0476
		ㄴ test_ref
		ㄴ train
		ㄴ train_ref
		ㄴ val
		ㄴ val_ref
	ㄴ night
    ㄴ rain
    ㄴ snow

 

따라서 카메라 이름과 시퀀스 번호가 일치해야 같은 위치에서 찍은 사진입니다. 

 

PyTorch Dataset Preparation

 

그러면 PyTorch 에서 이 데이터셋을 사용하기 위한 준비를 해 봅시다. 

우선 작업하는 폴더 내부에 <acdc_dataset_module> 폴더를 만들고 이 폴더 내부에 __init__.py 파일을 만들어줍니다. 

ㄴ workspace
	ㄴ acdc_dataset_module
    	ㄴ __init__.py 
    ㄴ temp.ipynb # 테스트할 ipynb 파일

 

ACDC 데이터셋으로 제가 만들고 싶은 데이터셋은 <adverse condition, good condition> 의 쌍들입니다. 

 

__init__.py 

 

import os
import torch
import pickle
from torch.utils.data import Dataset
from PIL import Image
import matplotlib.pyplot as plt

class ACDCDataset(Dataset):
    def __init__(self, root_dir, transform=None, rebuild=False):
        self.root_dir = root_dir
        self.transform = transform
        
        # pickle 을 이용해 객체 자체를 저장
        self.pickle_file = os.path.join(root_dir, "data_list.pkl")
        

        if rebuild or not os.path.exists(self.pickle_file):
            self.data_list = self._build_dataset()
            with open(self.pickle_file, 'wb') as f:
                pickle.dump(self.data_list, f)
        else:
            with open(self.pickle_file, 'rb') as f:
                self.data_list = pickle.load(f)
                
    def __getitem__(self, index):
        adverse_img_path, ref_img_path = self.data_list[index]
        adverse_img = Image.open(adverse_img_path)
        ref_img = Image.open(ref_img_path)

        if self.transform:
            adverse_img = self.transform(adverse_img)
            ref_img = self.transform(ref_img)

        return adverse_img, ref_img

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


    def _build_dataset(self):
        data_list = []
        conditions = ['fog', 'night', 'rain', 'snow']
        modes = ['train', 'val', 'test']

        for condition in conditions:
            for mode in modes:
                folder_path = os.path.join(self.root_dir, condition, mode)
                ref_folder_path = os.path.join(self.root_dir, condition, mode + '_ref')

                locations = os.listdir(folder_path)
                for location in locations:
                    images = os.listdir(os.path.join(folder_path, location))
                    for img_name in images:
                        img_path = os.path.join(folder_path, location, img_name)
                        ref_img_path = os.path.join(ref_folder_path, location, img_name.replace('rgb_anon', 'rgb_ref_anon'))
                        
                        # 이미지 쌍 저장
                        data_list.append((img_path, ref_img_path))
        return data_list

 

temp.ipynb

여기서 객체를 호출하고 잘 저장되었는지 확인합니다. 

 

from acdc_dataset_module import ACDCDataset
from torchvision import transforms
import matplotlib.pyplot as plt

# 데이터셋 초기화
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])
dataset = ACDCDataset("/material/data/acdc_dataset/rgb_anon_trainvaltest/rgb_anon", transform=transform)

 

dataset 의 길이를 확인해보면 다음과 같이 4006 쌍의 데이터셋이 잘 저장된 것을 볼 수 있습니다. 

# 처음 5개의 이미지 쌍을 시각화
for i in range(5):
    img, ref_img = dataset[i]
    
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.imshow(img.permute(1, 2, 0))
    plt.title(f"Image {i+1} - Adverse Condition")
    
    plt.subplot(1, 2, 2)
    plt.imshow(ref_img.permute(1, 2, 0))
    plt.title(f"Image {i+1} - Good Condition")
    
    plt.tight_layout()
    plt.show()

이와 같이 같은 장소에 대해서 각각 다른 condition 의 사진들이 리스트로 잘 저장된 것을 볼 수 있습니다.