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 VKITTIDataset(BaseDataset): | |
| def __init__(self, cfg, phase, **kwargs): | |
| super(VKITTIDataset, self).__init__( | |
| cfg=cfg, | |
| phase=phase, | |
| **kwargs) | |
| self.metric_scale = cfg.metric_scale | |
| def process_depth(self, depth, rgb): | |
| depth[depth>(150 * self.metric_scale)] = 0 | |
| depth /= self.metric_scale | |
| return depth | |
| def load_sem_label(self, sem_path, depth=None, sky_id=142) -> np.array: | |
| """ | |
| Category r g b | |
| Terrain 210 0 200 | |
| Sky 90 200 255 | |
| Tree 0 199 0 | |
| Vegetation 90 240 0 | |
| Building 140 140 140 | |
| Road 100 60 100 | |
| GuardRail 250 100 255 | |
| TrafficSign 255 255 0 | |
| TrafficLight 200 200 0 | |
| Pole 255 130 0 | |
| Misc 80 80 80 | |
| Truck 160 60 60 | |
| Car 255 127 80 | |
| Van 0 139 139 | |
| """ | |
| H, W = depth.shape | |
| sem_label = np.ones((H, W), dtype=np.int) * -1 | |
| sem = cv2.imread(sem_path)[:, :, ::-1] | |
| if sem is None: | |
| return sem_label | |
| sky_color = [90, 200, 255] | |
| sky_mask = (sem == sky_color).all(axis=2) | |
| sem_label[sky_mask] = 142 # set sky region to 142 | |
| return sem_label | |
| if __name__ == '__main__': | |
| from mmcv.utils import Config | |
| cfg = Config.fromfile('mono/configs/Apolloscape_DDAD/convnext_base.cascade.1m.sgd.mae.py') | |
| dataset_i = ApolloscapeDataset(cfg['Apolloscape'], 'train', **cfg.data_basic) | |
| print(dataset_i) | |