openpoints 0.1.0__py3-none-any.whl

openpoints/dataset/s3dis/s3dis_block.py

@@ -0,0 +1,96 @@
+import torch
+import torch.utils.data as data
+
+
+
+
+
+
+
+
+# class S3DIS(data.Dataset):
+#     r"""The (pre-processed) Stanford Large-Scale 3D Indoor Spaces dataset from
+#     the `"3D Semantic Parsing of Large-Scale Indoor Spaces"
+#     <https://openaccess.thecvf.com/content_cvpr_2016/papers/Armeni_3D_Semantic_Parsing_CVPR_2016_paper.pdf>`_
+#     paper, containing point clouds of six large-scale indoor parts in three
+#     buildings with 12 semantic elements (and one clutter class).
+
+#     Args:
+#         root (string): Root directory where the dataset should be saved.
+#         test_area (int, optional): Which area to use for testing (1-6).
+#             (default: :obj:`6`)
+#         train (bool, optional): If :obj:`True`, loads the training dataset,
+#             otherwise the test dataset. (default: :obj:`True`)
+#         transform (callable, optional): A function/transform that takes in an
+#             :obj:`torch_geometric.data.Data` object and returns a transformed
+#             version. The data object will be transformed before every access.
+#             (default: :obj:`None`)
+#         pre_transform (callable, optional): A function/transform that takes in
+#             an :obj:`torch_geometric.data.Data` object and returns a
+#             transformed version. The data object will be transformed before
+#             being saved to disk. (default: :obj:`None`)
+#         pre_filter (callable, optional): A function that takes in an
+#             :obj:`torch_geometric.data.Data` object and returns a boolean
+#             value, indicating whether the data object should be included in the
+#             final dataset. (default: :obj:`None`)
+#     """
+
+#     url = ('https://shapenet.cs.stanford.edu/media/'
+#            'indoor3d_sem_seg_hdf5_data.zip')
+
+#     def __init__(self, root, test_area=6, train=True, transform=None,
+#                  pre_transform=None, pre_filter=None):
+#         assert test_area >= 1 and test_area <= 6
+#         self.test_area = test_area
+#         super().__init__(root, transform, pre_transform, pre_filter)
+#         path = self.processed_paths[0] if train else self.processed_paths[1]
+#         self.data, self.slices = torch.load(path)
+
+#     @property
+#     def raw_file_names(self):
+#         return ['all_files.txt', 'room_filelist.txt']
+
+#     @property
+#     def processed_file_names(self):
+#         return [f'{split}_{self.test_area}.pt' for split in ['train', 'test']]
+
+#     def download(self):
+#         path = download_url(self.url, self.root)
+#         extract_zip(path, self.root)
+#         os.unlink(path)
+#         shutil.rmtree(self.raw_dir)
+#         name = self.url.split('/')[-1].split('.')[0]
+#         os.rename(osp.join(self.root, name), self.raw_dir)
+
+#     def process(self):
+#         import h5py
+
+#         with open(self.raw_paths[0], 'r') as f:
+#             filenames = [x.split('/')[-1] for x in f.read().split('\n')[:-1]]
+
+#         with open(self.raw_paths[1], 'r') as f:
+#             rooms = f.read().split('\n')[:-1]
+
+#         xs, ys = [], []
+#         for filename in filenames:
+#             f = h5py.File(osp.join(self.raw_dir, filename))
+#             xs += torch.from_numpy(f['data'][:]).unbind(0)
+#             ys += torch.from_numpy(f['label'][:]).to(torch.long).unbind(0)
+
+#         test_area = f'Area_{self.test_area}'
+#         train_data_list, test_data_list = [], []
+#         for i, (x, y) in enumerate(zip(xs, ys)):
+#             data = Data(pos=x[:, :3], x=x[:, 3:], y=y)
+
+#             if self.pre_filter is not None and not self.pre_filter(data):
+#                 continue
+#             if self.pre_transform is not None:
+#                 data = self.pre_transform(data)
+
+#             if test_area not in rooms[i]:
+#                 train_data_list.append(data)
+#             else:
+#                 test_data_list.append(data)
+
+#         torch.save(self.collate(train_data_list), self.processed_paths[0])
+#         torch.save(self.collate(test_data_list), self.processed_paths[1])

openpoints/dataset/s3dis/s3dis_sphere.py

@@ -0,0 +1,349 @@
+import torch
+import torch.utils.data as data
+import numpy as np
+import os
+import pickle
+from sklearn.neighbors import KDTree
+from ..build import DATASETS
+from ..data_util import crop_pc, voxelize
+from tqdm import tqdm
+import logging
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+@DATASETS.register_module()
+class S3DISSphere(data.Dataset):
+    label_to_names = {0: 'ceiling',
+                      1: 'floor',
+                      2: 'wall',
+                      3: 'beam',
+                      4: 'column',
+                      5: 'window',
+                      6: 'door',
+                      7: 'chair',
+                      8: 'table',
+                      9: 'bookcase',
+                      10: 'sofa',
+                      11: 'board',
+                      12: 'clutter'}
+    name_to_label = {v: i for i, v in enumerate(label_to_names.values())}
+    classes = list(label_to_names.values())
+    color_mean = np.array([0.5136457, 0.49523646, 0.44921124])
+    color_std = np.array([0.18308958, 0.18415008, 0.19252081])
+    num_classes = 13
+    class2color = {'ceiling':     [0, 255, 0],
+                   'floor':       [0, 0, 255],
+                   'wall':        [0, 255, 255],
+                   'beam':        [255, 255, 0],
+                   'column':      [255, 0, 255],
+                   'window':      [100, 100, 255],
+                   'door':        [200, 200, 100],
+                   'chair':       [255, 0, 0],
+                   'table':       [170, 120, 200],
+                   'bookcase':    [10, 200, 100],
+                   'sofa':        [200, 100, 100],
+                   'board':       [200, 200, 200],
+                   'clutter':     [50, 50, 50]}
+    cmap = np.array([*class2color.values()]).astype(np.uint8)
+    gravity_dim = 2
+    """S3DIS dataset for scene segmentation task.
+    Args:
+        voxel_size: grid length for pre-subsampling point clouds.
+        in_radius: radius of each input spheres.
+        num_points: max number of points for the input spheres.
+        num_steps: number of spheres for one training epoch.
+        num_epochs: total epochs.
+        data_root: root path for data.
+        transform: data transformations.
+        split: dataset split name.
+    """
+    def __init__(self, voxel_size,
+                 in_radius, num_points, num_steps, num_epochs,
+                 data_root=None,
+                 transform=None,
+                 split='train',
+                 centering=False,
+                 **kwargs):
+
+        super().__init__()
+        self.epoch = 0
+        self.transform = transform
+        self.voxel_size = voxel_size
+        self.in_radius = in_radius
+        self.num_points = num_points
+        self.num_steps = num_steps
+        self.num_epochs = num_epochs
+        self.centering = centering
+        self.data_root = data_root
+        self.train_clouds = ['Area_1', 'Area_2', 'Area_3', 'Area_4', 'Area_6']
+        self.val_clouds = ['Area_5']
+        self.split = split
+
+        if split == 'train':
+            cloud_names = self.train_clouds
+        else:
+            cloud_names = self.val_clouds
+
+        processed_dir = os.path.join(data_root, 'processed')
+        if not os.path.exists(processed_dir):
+            os.makedirs(processed_dir)
+
+        # prepare data
+        filename = os.path.join(
+            processed_dir, f'{split}_{voxel_size:.3f}_data.pkl')
+        if not os.path.exists(filename):
+            cloud_points_list, cloud_points_color_list, cloud_points_label_list = [], [], []
+            sub_cloud_points_list, sub_cloud_points_label_list, sub_cloud_points_color_list = [], [], []
+            sub_cloud_tree_list = []
+            cloud_rooms_list = []
+            for cloud_idx, cloud_name in enumerate(cloud_names):
+                # Pass if the cloud has already been computed
+                # Get rooms of the current cloud
+                cloud_folder = os.path.join(data_root,cloud_name)
+                room_folders = [os.path.join(cloud_folder, room) for room in os.listdir(cloud_folder) if
+                                os.path.isdir(os.path.join(cloud_folder, room))]
+                # Initiate containers
+                cloud_points = np.empty((0, 3), dtype=np.float32)
+                cloud_colors = np.empty((0, 3), dtype=np.float32)
+                cloud_classes = np.empty((0, 1), dtype=np.int32)
+                cloud_room_split = [0]
+                # Loop over rooms
+                for i, room_folder in enumerate(room_folders):
+                    print(
+                        'Cloud %s - Room %d/%d : %s' % (
+                            cloud_name, i + 1, len(room_folders), room_folder.split('\\')[-1]))
+                    room_npoints = []
+                    for object_name in os.listdir(os.path.join(room_folder, 'Annotations')):
+                        if object_name[-4:] == '.txt':
+                            # Text file containing point of the object
+                            object_file = os.path.join(room_folder, 'Annotations', object_name)
+                            # Object class and ID
+                            tmp = object_name[:-4].split('_')[0]
+                            if tmp in self.name_to_label:
+                                object_class = self.name_to_label[tmp]
+                            elif tmp in ['stairs']:
+                                object_class = self.name_to_label['clutter']
+                            else:
+                                raise ValueError(
+                                    'Unknown object name: ' + str(tmp))
+                            # Read object points and colors
+                            try:
+                                object_data = np.loadtxt(object_file)
+                            except:
+                                logging.info(f"error in reading file {object_file}")
+                            # Stack all data
+                            room_npoints.append(len(object_data))
+                            cloud_points = np.vstack(
+                                (cloud_points, object_data[:, 0:3].astype(np.float32)))
+                            cloud_colors = np.vstack(
+                                (cloud_colors, object_data[:, 3:6].astype(np.uint8)))
+                            object_classes = np.full(
+                                (object_data.shape[0], 1), object_class, dtype=np.int32)
+                            cloud_classes = np.vstack(
+                                (cloud_classes, object_classes))
+                    cloud_room_split.append(cloud_room_split[-1]+sum(room_npoints))
+                cloud_points_list.append(cloud_points)
+                cloud_points_color_list.append(cloud_colors)
+                cloud_points_label_list.append(cloud_classes)
+                cloud_rooms_list.append(cloud_room_split)
+                
+                sub_cloud_file = os.path.join(
+                    processed_dir, cloud_name + f'_{voxel_size:.3f}_sub.pkl')
+                if os.path.exists(sub_cloud_file):
+                    with open(sub_cloud_file, 'rb') as f:
+                        sub_points, sub_colors, sub_labels, search_tree = pickle.load(
+                            f)
+                else:
+                    if voxel_size > 0:
+                        sub_points, sub_colors, sub_labels = crop_pc(
+                            cloud_points, cloud_colors, cloud_classes, voxel_size=voxel_size)
+                        sub_labels = np.squeeze(sub_labels)
+                    else:
+                        sub_points = cloud_points
+                        sub_colors = cloud_colors
+                        sub_labels = cloud_classes
+
+                    # Get chosen neighborhoods
+                    search_tree = KDTree(sub_points, leaf_size=50)
+
+                    with open(sub_cloud_file, 'wb') as f:
+                        pickle.dump((sub_points, sub_colors,
+                                    sub_labels, search_tree), f)
+
+                sub_cloud_points_list.append(sub_points)
+                sub_cloud_points_color_list.append(sub_colors)
+                sub_cloud_points_label_list.append(sub_labels)
+                sub_cloud_tree_list.append(search_tree)
+
+            # original points
+            self.clouds_points = cloud_points_list
+            self.clouds_points_colors = cloud_points_color_list
+            self.clouds_points_labels = cloud_points_label_list
+            self.clouds_rooms = cloud_rooms_list
+            
+            # grid subsampled points
+            self.sub_clouds_points = sub_cloud_points_list
+            self.sub_clouds_points_colors = sub_cloud_points_color_list
+            self.sub_clouds_points_labels = sub_cloud_points_label_list
+            self.sub_cloud_trees = sub_cloud_tree_list
+
+            with open(filename, 'wb') as f:
+                pickle.dump((self.clouds_points, self.clouds_points_colors, self.clouds_points_labels, self.clouds_rooms, 
+                             self.sub_clouds_points, self.sub_clouds_points_colors, self.sub_clouds_points_labels,
+                             self.sub_cloud_trees), f)
+                print(f"{filename} saved successfully")
+        else:
+            with open(filename, 'rb') as f:
+                (self.clouds_points, self.clouds_points_colors, self.clouds_points_labels, 
+                 self.clouds_rooms, 
+                 self.sub_clouds_points, self.sub_clouds_points_colors, self.sub_clouds_points_labels,
+                 self.sub_cloud_trees) = pickle.load(f)
+                print(f"{filename} loaded successfully")
+
+        # prepare iteration indices
+        filename = os.path.join(processed_dir,
+                                f'{split}_{voxel_size:.3f}_{self.num_epochs}_{self.num_steps}_iterinds.pkl')
+        if not os.path.exists(filename):
+            potentials = []
+            min_potentials = []  # self.sub_cloud_trees, the tree for the grid_subsampled area
+            for cloud_i, tree in enumerate(self.sub_cloud_trees):
+                print(f"{split}/{cloud_i} has {tree.data.shape[0]} points")
+                cur_potential = np.random.rand(tree.data.shape[0]) * 1e-3
+                potentials.append(cur_potential)
+                min_potentials.append(float(np.min(cur_potential)))
+            self.cloud_inds = []
+            self.point_inds = []
+            self.noise = []
+            for ep in tqdm(range(self.num_epochs)):
+                for st in range(self.num_steps):
+                    # cloud_ind is the index for area in each split. for test, only 1 cloud.
+                    cloud_ind = int(np.argmin(min_potentials))
+                    point_ind = np.argmin(potentials[cloud_ind])
+                    self.cloud_inds.append(cloud_ind)
+                    self.point_inds.append(point_ind)
+                    points = np.array(
+                        self.sub_cloud_trees[cloud_ind].data, copy=False)  # sub_points
+                    center_point = points[point_ind, :].reshape(1, -1)
+                    noise = np.random.normal(
+                        scale=self.in_radius / 10, size=center_point.shape)
+                    self.noise.append(noise)
+                    pick_point = center_point + noise.astype(center_point.dtype)
+                    # Indices of points in input region
+                    query_inds = self.sub_cloud_trees[cloud_ind].query_radius(pick_point,
+                                                                              r=self.in_radius,
+                                                                              return_distance=True,
+                                                                              sort_results=True)[0][0]
+                    cur_num_points = query_inds.shape[0]
+                    if self.num_points < cur_num_points:
+                        query_inds = query_inds[:self.num_points]
+                    # Update potentials (Tuckey weights)
+                    dists = np.sum(
+                        np.square((points[query_inds] - pick_point).astype(np.float32)), axis=1)
+                    tukeys = np.square(1 - dists / np.square(self.in_radius))
+                    tukeys[dists > np.square(self.in_radius)] = 0
+                    potentials[cloud_ind][query_inds] += tukeys
+                    min_potentials[cloud_ind] = float(
+                        np.min(potentials[cloud_ind]))
+            with open(filename, 'wb') as f:
+                pickle.dump((self.cloud_inds, self.point_inds, self.noise), f)
+                print(f"{filename} saved successfully")
+        else:
+            with open(filename, 'rb') as f:
+                self.cloud_inds, self.point_inds, self.noise = pickle.load(f)
+                print(f"{filename} loaded successfully")
+
+        # prepare validation projection inds
+        filename = os.path.join(
+            processed_dir, f'{split}_{voxel_size:.3f}_proj.pkl')
+        if not os.path.exists(filename):
+            proj_ind_list = []
+            for points, search_tree in zip(self.clouds_points, self.sub_cloud_trees):
+                # points: the original points.  700K points
+                # subcloud_trees: the tree of the subsampeld points. 500K points
+                # find the nearest point in the subsampled points for each point in the original points.
+                proj_inds = np.squeeze(search_tree.query(
+                    points, return_distance=False))
+                proj_inds = proj_inds.astype(np.int32)
+                proj_ind_list.append(proj_inds)
+            self.projections = proj_ind_list
+            with open(filename, 'wb') as f:
+                pickle.dump(self.projections, f)
+                print(f"{filename} saved successfully")
+        else:
+            with open(filename, 'rb') as f:
+                self.projections = pickle.load(f)
+                print(f"{filename} loaded successfully")
+
+    def __getitem__(self, idx):
+        """
+        Returns:
+            pts: (N, 3), a point cloud.
+            mask: (N, ), 0/1 mask to distinguish padding points.
+            features: (input_features_dim, N), input points features.
+            pts_labels: (N), point label.
+            current_cloud_index: (1), cloud index.
+            input_inds: (N), the index of input points in point cloud.
+        """
+        cloud_ind = self.cloud_inds[idx + self.epoch * self.num_steps]
+        point_ind = self.point_inds[idx + self.epoch * self.num_steps]
+        noise = self.noise[idx + self.epoch * self.num_steps]
+        points = np.array(
+            self.sub_cloud_trees[cloud_ind].data, copy=False)  # subpoints
+        center_point = points[point_ind, :].reshape(1, -1)
+        pick_point = center_point + noise.astype(center_point.dtype)
+        # Indices of points in input region
+        # the point index in the subsampled point tree.
+        query_inds = self.sub_cloud_trees[cloud_ind].query_radius(pick_point,
+                                                                  r=self.in_radius,
+                                                                  return_distance=True,
+                                                                  sort_results=True)[0][0]
+        # Number collected
+        cur_num_points = query_inds.shape[0]
+        if self.num_points < cur_num_points:
+            shuffle_choice = np.random.permutation(np.arange(self.num_points))
+            input_inds = query_inds[:self.num_points][shuffle_choice]
+            mask = torch.ones(self.num_points).type(torch.int32)
+        else:
+            shuffle_choice = np.random.permutation(np.arange(cur_num_points))
+            query_inds = query_inds[shuffle_choice]
+            padding_choice = np.random.choice(
+                cur_num_points, self.num_points - cur_num_points)
+            input_inds = np.hstack([query_inds, query_inds[padding_choice]])
+            mask = torch.zeros(self.num_points).type(torch.int32)
+            mask[:cur_num_points] = 1
+        # points: the whole point of an area (voxel downsampled). 
+        # input_inds: the point index of the current area.
+        original_points = points[input_inds]
+        pts = (original_points - pick_point).astype(np.float32)
+        colors = self.sub_clouds_points_colors[cloud_ind][input_inds].astype(
+            np.float32)
+        labels = self.sub_clouds_points_labels[cloud_ind][input_inds].astype(
+            np.int64)
+        current_cloud_index = np.array(cloud_ind).astype(np.int64)
+
+        data = {'pos': pts,
+                'x': colors,
+                'y': labels,
+                'mask': mask,
+                'cloud_index': current_cloud_index,
+                'input_inds': input_inds,
+                }
+        """ vis
+        from openpoints.dataset import vis_multi_points, vis_points
+        vis_multi_points([points, pts], [self.sub_clouds_points_colors[cloud_ind]/255., colors/255.])
+        
+        import copy
+        old_data = copy.deepcopy(data)
+        if self.transform is not None:
+            data = self.transform(data)
+        vis_multi_points([old_data['pos'][:, :3], data['pos'][:, :3].numpy()], colors=[old_data['x'][:, :3]/255.,data['x'][:, :3].numpy()])
+        """   
+        if self.transform is not None:
+            data = self.transform(data)
+
+        if 'heights' not in data.keys():
+            data['heights'] =  torch.from_numpy(original_points[:, self.gravity_dim:self.gravity_dim+1].astype(np.float32))
+        return data
+
+    def __len__(self):
+        return self.num_steps

openpoints/dataset/scannetv2/__init__.py

@@ -0,0 +1 @@
+from .scannet import ScanNet

openpoints/dataset/scannetv2/scannet.py

@@ -0,0 +1,176 @@
+import os
+import os.path as osp
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from ..build import DATASETS
+from ..data_util import crop_pc, voxelize
+from ...transforms.point_transform_cpu import PointsToTensor
+import glob
+from tqdm import tqdm
+import logging
+import pickle
+
+
+VALID_CLASS_IDS = [
+    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39
+]
+
+SCANNET_COLOR_MAP = {
+    0: (0., 0., 0.),
+    1: (174., 199., 232.),
+    2: (152., 223., 138.),
+    3: (31., 119., 180.),
+    4: (255., 187., 120.),
+    5: (188., 189., 34.),
+    6: (140., 86., 75.),
+    7: (255., 152., 150.),
+    8: (214., 39., 40.),
+    9: (197., 176., 213.),
+    10: (148., 103., 189.),
+    11: (196., 156., 148.),
+    12: (23., 190., 207.),
+    14: (247., 182., 210.),
+    15: (66., 188., 102.),
+    16: (219., 219., 141.),
+    17: (140., 57., 197.),
+    18: (202., 185., 52.),
+    19: (51., 176., 203.),
+    20: (200., 54., 131.),
+    21: (92., 193., 61.),
+    22: (78., 71., 183.),
+    23: (172., 114., 82.),
+    24: (255., 127., 14.),
+    25: (91., 163., 138.),
+    26: (153., 98., 156.),
+    27: (140., 153., 101.),
+    28: (158., 218., 229.),
+    29: (100., 125., 154.),
+    30: (178., 127., 135.),
+    32: (146., 111., 194.),
+    33: (44., 160., 44.),
+    34: (112., 128., 144.),
+    35: (96., 207., 209.),
+    36: (227., 119., 194.),
+    37: (213., 92., 176.),
+    38: (94., 106., 211.),
+    39: (82., 84., 163.),
+    40: (100., 85., 144.),
+}
+
+
+
+@DATASETS.register_module()
+class ScanNet(Dataset):
+    num_classes = 20
+    classes = ['wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window', 'bookshelf', 'picture',
+               'counter', 'desk', 'curtain', 'refrigerator', 'shower curtain', 'toilet', 'sink', 'bathtub', 'otherfurniture']
+    gravity_dim = 2
+    
+    color_mean = [0.46259782, 0.46253258, 0.46253258]
+    color_std =  [0.693565  , 0.6852543 , 0.68061745]
+    """ScanNet dataset, loading the subsampled entire room as input without block/sphere subsampling.
+    number of points per room in average, median, and std: (145841.0, 158783.87179487178, 84200.84445829492)
+    """  
+    def __init__(self,
+                 data_root='data/ScanNet',
+                 split='train',
+                 voxel_size=0.04,
+                 voxel_max=None,
+                 transform=None,
+                 loop=1, presample=False, variable=False,
+                 n_shifted=1
+                 ):
+        super().__init__()
+        self.split = split
+        self.voxel_size = voxel_size
+        self.voxel_max = voxel_max
+        self.transform = transform
+        self.presample = presample
+        self.variable = variable
+        self.loop = loop
+        self.n_shifted = n_shifted
+        self.pipe_transform = PointsToTensor() 
+
+        if split == "train" or split == 'val':
+            self.data_list = glob.glob(os.path.join(data_root, split, "*.pth"))
+        elif split == 'trainval':
+            self.data_list = glob.glob(os.path.join(
+                data_root, "train", "*.pth")) + glob.glob(os.path.join(data_root, "val", "*.pth"))
+        elif split == 'test':
+            self.data_list = glob.glob(os.path.join(data_root, split, "*.pth"))
+        else:
+            raise ValueError("no such split: {}".format(split))
+
+        logging.info("Totally {} samples in {} set.".format(
+            len(self.data_list), split))
+
+        processed_root = os.path.join(data_root, 'processed')
+        filename = os.path.join(
+            processed_root, f'scannet_{split}_{voxel_size:.3f}.pkl')
+        if presample and not os.path.exists(filename):
+            np.random.seed(0)
+            self.data = []
+            for item in tqdm(self.data_list, desc=f'Loading ScanNet {split} split'):
+                data = torch.load(item)
+                coord, feat, label = data[0:3]
+                coord, feat, label = crop_pc(
+                    coord, feat, label, self.split, self.voxel_size, self.voxel_max, variable=self.variable)
+                cdata = np.hstack(
+                    (coord, feat, np.expand_dims(label, -1))).astype(np.float32)
+                self.data.append(cdata)
+            npoints = np.array([len(data) for data in self.data])
+            logging.info('split: %s, median npoints %.1f, avg num points %.1f, std %.1f' % (
+                self.split, np.median(npoints), np.average(npoints), np.std(npoints)))
+            os.makedirs(processed_root, exist_ok=True)
+            with open(filename, 'wb') as f:
+                pickle.dump(self.data, f)
+                print(f"{filename} saved successfully")
+        elif presample:
+            with open(filename, 'rb') as f:
+                self.data = pickle.load(f)
+                print(f"{filename} load successfully")
+            # median, average, std of number of points after voxel sampling for val set.
+            # (100338.5, 109686.1282051282, 57024.51083415437)
+            # before voxel sampling
+            # (145841.0, 158783.87179487178, 84200.84445829492)
+    def __getitem__(self, idx):
+        data_idx = idx % len(self.data_list)
+        if self.presample:
+            coord, feat, label = np.split(self.data[data_idx], [3, 6], axis=1)
+        else:
+            data_path = self.data_list[data_idx]
+            data = torch.load(data_path)
+            coord, feat, label = data[0:3]
+
+        feat = (feat + 1) * 127.5
+        label = label.astype(np.long).squeeze()
+        data = {'pos': coord.astype(np.float32), 'x': feat.astype(np.float32), 'y': label}
+        """debug 
+        from openpoints.dataset import vis_multi_points
+        import copy
+        old_data = copy.deepcopy(data)
+        if self.transform is not None:
+            data = self.transform(data)
+        data['pos'], data['x'], data['y'] = crop_pc(
+            data['pos'], data['x'], data['y'], self.split, self.voxel_size, self.voxel_max,
+            downsample=not self.presample, variable=self.variable)
+            
+        vis_multi_points([old_data['pos'][:, :3], data['pos'][:, :3]], colors=[old_data['x'][:, :3]/255.,data['x'][:, :3]])
+        """
+        if self.transform is not None:
+            data = self.transform(data)
+        
+        if not self.presample: 
+            data['pos'], data['x'], data['y'] = crop_pc(
+                data['pos'], data['x'], data['y'], self.split, self.voxel_size, self.voxel_max,
+                downsample=not self.presample, variable=self.variable)
+        
+        data = self.pipe_transform(data)
+         
+        if 'heights' not in data.keys():
+            data['heights'] =  data['pos'][:, self.gravity_dim:self.gravity_dim+1] - data['pos'][:, self.gravity_dim:self.gravity_dim+1].min()
+        return data
+
+    def __len__(self):
+        return len(self.data_list) * self.loop

openpoints/dataset/scanobjectnn/__init__.py

@@ -0,0 +1,3 @@
+from .scanobjectnn import ScanObjectNNHardest
+
+