Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,117 Bytes
4893ce0 |
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 |
"""
Structured3D Datasets
Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""
import os
import glob
from collections.abc import Sequence
from .defaults import DefaultDataset
from .builder import DATASETS
@DATASETS.register_module()
class Structured3DDataset(DefaultDataset):
def get_data_list(self):
if isinstance(self.split, str):
data_list = glob.glob(
os.path.join(self.data_root, self.split, "scene_*/room_*")
)
elif isinstance(self.split, Sequence):
data_list = []
for split in self.split:
data_list += glob.glob(
os.path.join(self.data_root, split, "scene_*/room_*")
)
else:
raise NotImplementedError
return data_list
def get_data_name(self, idx):
file_path = self.data_list[idx % len(self.data_list)]
dir_path, room_name = os.path.split(file_path)
scene_name = os.path.basename(dir_path)
data_name = f"{scene_name}_{room_name}"
return data_name
|