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 UASOLDataset(BaseDataset): | |
| def __init__(self, cfg, phase, **kwargs): | |
| super(UASOLDataset, self).__init__( | |
| cfg=cfg, | |
| phase=phase, | |
| **kwargs) | |
| self.metric_scale = cfg.metric_scale | |
| def process_depth(self, depth, rgb): | |
| depth[depth>65500] = 0 | |
| depth /= self.metric_scale | |
| return depth | |
| def load_rgb_depth(self, rgb_path: str, depth_path: str) -> (np.array, np.array): | |
| """ | |
| Load the rgb and depth map with the paths. | |
| """ | |
| rgb = self.load_data(rgb_path, is_rgb_img=True) | |
| if rgb is None: | |
| self.logger.info(f'>>>>{rgb_path} has errors.') | |
| depth = self.load_data(depth_path) | |
| if depth is None: | |
| self.logger.info(f'{depth_path} has errors.') | |
| depth = depth.astype(np.float) | |
| depth = self.process_depth(depth, rgb) | |
| depth = depth[1:-1, ...] | |
| return rgb, depth | |
| if __name__ == '__main__': | |
| from mmcv.utils import Config | |
| cfg = Config.fromfile('mono/configs/Apolloscape_DDAD/convnext_base.cascade.1m.sgd.mae.py') | |
| dataset_i = UASOLDataset(cfg['Apolloscape'], 'train', **cfg.data_basic) | |
| print(dataset_i) | |