Spaces:
Running
Running
File size: 2,458 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import pickle
import os
import os.path as osp
# RGBD-Dataset
from .tartan import TartanAir
from .stream import ImageStream
from .stream import StereoStream
from .stream import RGBDStream
# streaming datasets for inference
from .tartan import TartanAirStream
from .tartan import TartanAirTestStream
def dataset_factory(dataset_list, **kwargs):
""" create a combined dataset """
from torch.utils.data import ConcatDataset
dataset_map = { 'tartan': (TartanAir, ) }
db_list = []
for key in dataset_list:
# cache datasets for faster future loading
db = dataset_map[key][0](**kwargs)
print("Dataset {} has {} images".format(key, len(db)))
db_list.append(db)
return ConcatDataset(db_list)
def create_datastream(dataset_path, **kwargs):
""" create data_loader to stream images 1 by 1 """
from torch.utils.data import DataLoader
if osp.isfile(osp.join(dataset_path, 'calibration.txt')):
db = ETH3DStream(dataset_path, **kwargs)
elif osp.isdir(osp.join(dataset_path, 'image_left')):
db = TartanAirStream(dataset_path, **kwargs)
elif osp.isfile(osp.join(dataset_path, 'rgb.txt')):
db = TUMStream(dataset_path, **kwargs)
elif osp.isdir(osp.join(dataset_path, 'mav0')):
db = EurocStream(dataset_path, **kwargs)
elif osp.isfile(osp.join(dataset_path, 'calib.txt')):
db = KITTIStream(dataset_path, **kwargs)
else:
# db = TartanAirStream(dataset_path, **kwargs)
db = TartanAirTestStream(dataset_path, **kwargs)
stream = DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
return stream
def create_imagestream(dataset_path, **kwargs):
""" create data_loader to stream images 1 by 1 """
from torch.utils.data import DataLoader
db = ImageStream(dataset_path, **kwargs)
return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
def create_stereostream(dataset_path, **kwargs):
""" create data_loader to stream images 1 by 1 """
from torch.utils.data import DataLoader
db = StereoStream(dataset_path, **kwargs)
return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
def create_rgbdstream(dataset_path, **kwargs):
""" create data_loader to stream images 1 by 1 """
from torch.utils.data import DataLoader
db = RGBDStream(dataset_path, **kwargs)
return DataLoader(db, shuffle=False, batch_size=1, num_workers=4)
|