| 68 | |
| 69 | |
| 70 | class DepthDataLoader(object): |
| 71 | def __init__(self, config, mode, device='cpu', transform=None, **kwargs): |
| 72 | """ |
| 73 | Data loader for depth datasets |
| 74 | |
| 75 | Args: |
| 76 | config (dict): Config dictionary. Refer to utils/config.py |
| 77 | mode (str): "train" or "online_eval" |
| 78 | device (str, optional): Device to load the data on. Defaults to 'cpu'. |
| 79 | transform (torchvision.transforms, optional): Transform to apply to the data. Defaults to None. |
| 80 | """ |
| 81 | |
| 82 | self.config = config |
| 83 | |
| 84 | if config.dataset == 'ibims': |
| 85 | self.data = get_ibims_loader(config, batch_size=1, num_workers=1) |
| 86 | return |
| 87 | |
| 88 | if config.dataset == 'sunrgbd': |
| 89 | self.data = get_sunrgbd_loader( |
| 90 | data_dir_root=config.sunrgbd_root, batch_size=1, num_workers=1) |
| 91 | return |
| 92 | |
| 93 | if config.dataset == 'diml_indoor': |
| 94 | self.data = get_diml_indoor_loader( |
| 95 | data_dir_root=config.diml_indoor_root, batch_size=1, num_workers=1) |
| 96 | return |
| 97 | |
| 98 | if config.dataset == 'diml_outdoor': |
| 99 | self.data = get_diml_outdoor_loader( |
| 100 | data_dir_root=config.diml_outdoor_root, batch_size=1, num_workers=1) |
| 101 | return |
| 102 | |
| 103 | if "diode" in config.dataset: |
| 104 | self.data = get_diode_loader( |
| 105 | config[config.dataset+"_root"], batch_size=1, num_workers=1) |
| 106 | return |
| 107 | |
| 108 | if config.dataset == 'hypersim_test': |
| 109 | self.data = get_hypersim_loader( |
| 110 | config.hypersim_test_root, batch_size=1, num_workers=1) |
| 111 | return |
| 112 | |
| 113 | if config.dataset == 'vkitti': |
| 114 | self.data = get_vkitti_loader( |
| 115 | config.vkitti_root, batch_size=1, num_workers=1) |
| 116 | return |
| 117 | |
| 118 | if config.dataset == 'vkitti2': |
| 119 | self.data = get_vkitti2_loader( |
| 120 | config.vkitti2_root, batch_size=1, num_workers=1) |
| 121 | return |
| 122 | |
| 123 | if config.dataset == 'ddad': |
| 124 | self.data = get_ddad_loader(config.ddad_root, resize_shape=( |
| 125 | 352, 1216), batch_size=1, num_workers=1) |
| 126 | return |
| 127 | |