Spaces:
Running
Running
| import os | |
| import json | |
| import torch | |
| import torchvision.transforms as transforms | |
| import os.path | |
| import numpy as np | |
| import cv2 | |
| from torch.utils.data import Dataset | |
| import random | |
| from .__base_dataset__ import BaseDataset | |
| class FisheyeDataset(BaseDataset): | |
| def __init__(self, cfg, phase, **kwargs): | |
| super(FisheyeDataset, self).__init__( | |
| cfg=cfg, | |
| phase=phase, | |
| **kwargs) | |
| self.metric_scale = cfg.metric_scale | |
| def load_data(self, path: str, is_rgb_img: bool=False): | |
| if not os.path.exists(path): | |
| self.logger.info(f'>>>>{path} does not exist.') | |
| # raise RuntimeError(f'{path} does not exist.') | |
| data_type = os.path.splitext(path)[-1] | |
| if data_type in self.img_file_type: | |
| if is_rgb_img: | |
| data = cv2.imread(path) | |
| else: | |
| data = cv2.imread(path, -1) | |
| data[data>65500] = 0 | |
| data &= 0x7FFF | |
| elif data_type in self.np_file_type: | |
| data = np.load(path) | |
| else: | |
| raise RuntimeError(f'{data_type} is not supported in current version.') | |
| return data.squeeze() | |
| def load_batch(self, meta_data, data_path): | |
| curr_intrinsic = meta_data['cam_in'] | |
| # load rgb/depth | |
| curr_rgb, curr_depth = self.load_rgb_depth(data_path['rgb_path'], data_path['depth_path']) | |
| # get semantic labels | |
| curr_sem = self.load_sem_label(data_path['sem_path'], curr_depth) | |
| # create camera model | |
| curr_cam_model = self.create_cam_model(curr_rgb.shape[0], curr_rgb.shape[1], curr_intrinsic) | |
| # get normal labels | |
| curr_normal = self.load_norm_label(data_path['normal_path'], H=curr_rgb.shape[0], W=curr_rgb.shape[1]) | |
| # get depth mask | |
| depth_mask = self.load_depth_valid_mask(data_path['depth_mask_path'])[:, :, :] | |
| # with masks from andy | |
| curr_depth[~(depth_mask[:, :, 0])] = -1 | |
| curr_rgb[~(depth_mask[:, :, :])] = 0 | |
| # get stereo depth | |
| curr_stereo_depth = self.load_stereo_depth_label(data_path['disp_path'], H=curr_rgb.shape[0], W=curr_rgb.shape[1]) | |
| data_batch = dict( | |
| curr_rgb = curr_rgb, | |
| curr_depth = curr_depth, | |
| curr_sem = curr_sem, | |
| curr_normal = curr_normal, | |
| curr_cam_model=curr_cam_model, | |
| curr_stereo_depth=curr_stereo_depth, | |
| ) | |
| return data_batch | |
| def process_depth(self, depth, rgb): | |
| depth /= self.metric_scale | |
| return depth | |