KITTI DataLoader 구현하기!

먼저 dataset.py를 생성합니다.

 

Data는 KITTI dataset이고, left image, left depth gt, left segmentation gt를 load할 것 입니다.

 

1. Import 및 label_classes 만들기

먼저 필요한 것들을 모두 import 해줍니다.

그리고 segmentation을 할 때 사람이 빨간색, 자동차가 파란색으로 표현이 되어, 사람 = (255,0,0), 자동차 = (0,0,255)라고 가정을 하면 해당하는 것을 labeling을 해주어야 한다. KITTI dataset은 19개의 label로 표현이 되며, label_classes는 위의 코드에 작성한 것대로이다. (ex, (0,0,0) = 255 - 검은색 부분으로 칠해진 것은 label이 255 여기서 255는 무시 되는 label로 0-18까지가 의미있는 label이다.)

 

그리고, 우리는 disparity data를 받아와서 depth로 바꾸어 줄 것이기 때문에 KITTI dataset의 focal length와 baseline을 작성해준다. 

 

2. KITTI_Dataloader class 생성, __init__ 함수 작성

KITTI_Dataloader class를 설정해준다. left와 disp, seg를 담을 빈 리스트를 생성해준다.

KITTI dataloader의 안에 init 함수에서 img와 mask disp의 경로를 설정하고, 위에서 설정한 left, disp, seg 빈 리스트에 이미지 경로를 넣어준다.

 

조금더 자세히 설명을 해보자면 left image가 모두 담겨져있는 directory를 "/home/sunkyung/Downloads/KITTI/image/train/left/"라고 한다면 이것을 self.img_root라고 한다.  그리고 나서, 아래의 for문을 통해서 img_root안에 있는 모든 이미지들을 읽어와서 앞의 경로와 함께 이미지 이름을 위에서 생성한 빈 리스트들에 넣어준다. (여기서 for문에 self.mask_root와 self.img_root, self.disp_root 안에 있는 이미지의 이름이 모두 동일하다. 따라서 os.listdir(self.mask_root)나 os.listdir(self.img_root)이거나 상관이 없다. ) "/home/sunkyung/Downloads/KITTI/image/train/left/000001.png" 이런식으로 directory 안에 있는 모든 이미지들이 경로와 함께 각 __left, __disp, __seg 리스트 안에 들어가게 될 것이다. 

 

3. KITTI_Dataloader class 안의 __getitem__ 함수 구현

 

Class KITTI_Dataloader 안에 __getitem__이라는 함수로 실제 이미지들을 가져와 이미지 데이터 처리를 해줄 것이다.

이미지 처리는 이미지 읽기, 사이즈 조절, np.array로 바꾸어 주기, normalize 하기, tensor로 바꾸기 이게 다라고 생각하면 된다. 그리고 augmentation을 하고 싶다면, crop, resize, flip, 색상 변환 등 다양하게 줄 수도 있다. 

 

먼저, Image.open함수를 사용하여 경로를 넣으면 image를 return 해주도록 한다. 그렇게 image를 얻어 온 후 먼저, img1은 left input image로 RGB 이미지이다. 이것은 전처리로, get_transform 함수( 이미지를 resize, normalize, tensor로 변환) 라는 것을 만들어 img1을 넣어준다. 그리고, img2는 disparity image로 disparity를 depth로 바꾸고, 사이즈를 변환하기 위해서 load_image2라는 함수를 만들어서 return 해준다. 그리고 img3는 segmentation image로 위에서 설정한 label_classes를 가지고 labeling을 해준다. 

 

위에서 구현한 KITTI_Dataloader()를 train에서 가지고 올 때 위와 같이 코드를 작성해서 사용하면 된다. 

 

이상 간단한 KITTI Dataloader 구현 코드 였습니다.

 

질문이 있다면 댓글로 남겨주세요.