|
import cv2 |
|
import torch |
|
from torch.utils.data import Dataset |
|
from torchvision.transforms import Compose |
|
|
|
from dataset.transform import Resize, NormalizeImage, PrepareForNet, Crop |
|
|
|
|
|
class VKITTI2(Dataset): |
|
def __init__(self, filelist_path, mode, size=(518, 518)): |
|
|
|
self.mode = mode |
|
self.size = size |
|
|
|
with open(filelist_path, 'r') as f: |
|
self.filelist = f.read().splitlines() |
|
|
|
net_w, net_h = size |
|
self.transform = Compose([ |
|
Resize( |
|
width=net_w, |
|
height=net_h, |
|
resize_target=True if mode == 'train' else False, |
|
keep_aspect_ratio=True, |
|
ensure_multiple_of=14, |
|
resize_method='lower_bound', |
|
image_interpolation_method=cv2.INTER_CUBIC, |
|
), |
|
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
|
PrepareForNet(), |
|
] + ([Crop(size[0])] if self.mode == 'train' else [])) |
|
|
|
def __getitem__(self, item): |
|
img_path = self.filelist[item].split(' ')[0] |
|
depth_path = self.filelist[item].split(' ')[1] |
|
|
|
image = cv2.imread(img_path) |
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0 |
|
|
|
depth = cv2.imread(depth_path, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) / 100.0 |
|
|
|
sample = self.transform({'image': image, 'depth': depth}) |
|
|
|
sample['image'] = torch.from_numpy(sample['image']) |
|
sample['depth'] = torch.from_numpy(sample['depth']) |
|
|
|
sample['valid_mask'] = (sample['depth'] <= 80) |
|
|
|
sample['image_path'] = self.filelist[item].split(' ')[0] |
|
|
|
return sample |
|
|
|
def __len__(self): |
|
return len(self.filelist) |
