32 lines
1.0 KiB
Python
32 lines
1.0 KiB
Python
import pickle
|
|
import numpy as np
|
|
|
|
from ._point_dataset import _Point_Dataset
|
|
|
|
|
|
class FullCloudsDataset(_Point_Dataset):
|
|
|
|
setting = 'grid'
|
|
|
|
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)
|
|
points = np.stack((pointcloud['x'], pointcloud['y'], pointcloud['z'],
|
|
pointcloud['xn'], pointcloud['yn'], pointcloud['zn']
|
|
),
|
|
axis=-1)
|
|
|
|
# When yopu want to return points and normal seperately
|
|
# normal = np.stack((pointcloud['xn'], pointcloud['yn'], pointcloud['zn']), axis=-1)
|
|
label = np.stack((pointcloud['label'], pointcloud['cl_idx']))
|
|
sample_idxs = self.sampling(points)
|
|
|
|
return points[sample_idxs], label[sample_idxs] |