|
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 |
|
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) |
|
|
