File size: 3,784 Bytes
57746f1 |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
"""
ArkitScenes Dataset
Author: Xiaoyang Wu ([email protected])
Please cite our work if the code is helpful to you.
"""
import os
import glob
import numpy as np
import torch
from copy import deepcopy
from torch.utils.data import Dataset
from pointcept.utils.logger import get_root_logger
from .builder import DATASETS
from .transform import Compose, TRANSFORMS
from .preprocessing.scannet.meta_data.scannet200_constants import VALID_CLASS_IDS_200
@DATASETS.register_module()
class ArkitScenesDataset(Dataset):
def __init__(
self,
split="Training",
data_root="data/ARKitScenesMesh",
transform=None,
test_mode=False,
test_cfg=None,
loop=1,
):
super(ArkitScenesDataset, self).__init__()
self.data_root = data_root
self.split = split
self.transform = Compose(transform)
self.loop = (
loop if not test_mode else 1
) # force make loop = 1 while in test mode
self.test_mode = test_mode
self.test_cfg = test_cfg if test_mode else None
self.class2id = np.array(VALID_CLASS_IDS_200)
if test_mode:
self.test_voxelize = TRANSFORMS.build(self.test_cfg.voxelize)
self.test_crop = TRANSFORMS.build(self.test_cfg.crop)
self.post_transform = Compose(self.test_cfg.post_transform)
self.aug_transform = [Compose(aug) for aug in self.test_cfg.aug_transform]
self.data_list = self.get_data_list()
logger = get_root_logger()
logger.info(
"Totally {} x {} samples in {} set.".format(
len(self.data_list), self.loop, split
)
)
def get_data_list(self):
if isinstance(self.split, str):
data_list = glob.glob(os.path.join(self.data_root, self.split, "*.pth"))
elif isinstance(self.split, list):
data_list = []
for split in self.split:
data_list += glob.glob(os.path.join(self.data_root, split, "*.pth"))
else:
raise NotImplementedError
return data_list
def get_data(self, idx):
data = torch.load(self.data_list[idx % len(self.data_list)])
coord = data["coord"]
color = data["color"]
normal = data["normal"]
segment = np.zeros(coord.shape[0])
data_dict = dict(coord=coord, normal=normal, color=color, segment=segment)
return data_dict
def get_data_name(self, idx):
data_idx = self.data_idx[idx % len(self.data_idx)]
return os.path.basename(self.data_list[data_idx]).split(".")[0]
def prepare_train_data(self, idx):
# load data
data_dict = self.get_data(idx)
data_dict = self.transform(data_dict)
return data_dict
def prepare_test_data(self, idx):
# load data
data_dict = self.get_data(idx)
segment = data_dict.pop("segment")
data_dict = self.transform(data_dict)
data_dict_list = []
for aug in self.aug_transform:
data_dict_list.append(aug(deepcopy(data_dict)))
input_dict_list = []
for data in data_dict_list:
data_part_list = self.test_voxelize(data)
for data_part in data_part_list:
data_part_list = self.test_crop(data_part)
input_dict_list += data_part_list
for i in range(len(input_dict_list)):
input_dict_list[i] = self.post_transform(input_dict_list[i])
return input_dict_list, segment
def __getitem__(self, idx):
if self.test_mode:
return self.prepare_test_data(idx)
else:
return self.prepare_train_data(idx)
def __len__(self):
return len(self.data_list) * self.loop
|