37 lines
990 B
Python
37 lines
990 B
Python
import pickle
|
|
from collections import defaultdict
|
|
|
|
import numpy as np
|
|
|
|
from ._point_dataset import _Point_Dataset
|
|
|
|
|
|
class FullCloudsDataset(_Point_Dataset):
|
|
|
|
setting = 'pc'
|
|
split: str
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super(FullCloudsDataset, self).__init__(*args, **kwargs)
|
|
|
|
def __len__(self):
|
|
return len(self._files)
|
|
|
|
def __getitem__(self, item):
|
|
processed_file_path = self._read_or_load(item)
|
|
|
|
with processed_file_path.open('rb') as processed_file:
|
|
pointcloud = pickle.load(processed_file)
|
|
|
|
position = np.stack((pointcloud['x'], pointcloud['y'], pointcloud['z']), axis=-1)
|
|
|
|
normal = np.stack((pointcloud['xn'], pointcloud['yn'], pointcloud['zn']), axis=-1)
|
|
|
|
label = pointcloud['label']
|
|
|
|
sample_idxs = self.sampling(position)
|
|
|
|
return (normal[sample_idxs].astype(np.float),
|
|
position[sample_idxs].astype(np.float),
|
|
label[sample_idxs].astype(np.int))
|