Spaces:
Running
Running
File size: 1,795 Bytes
b7eedf7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
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)
|