openpoints 0.1.0__py3-none-any.whl

openpoints/dataset/parsers/parser_tfds.py

@@ -0,0 +1,297 @@
+""" Dataset parser interface that wraps TFDS datasets
+
+Wraps many (most?) TFDS image-classification datasets
+from https://github.com/tensorflow/datasets
+https://www.tensorflow.org/datasets/catalog/overview#image_classification
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import math
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+try:
+    import tensorflow as tf
+    tf.config.set_visible_devices([], 'GPU')  # Hands off my GPU! (or pip install tensorflow-cpu)
+    import tensorflow_datasets as tfds
+    try:
+        tfds.even_splits('', 1, drop_remainder=False)  # non-buggy even_splits has drop_remainder arg
+        has_buggy_even_splits = False
+    except TypeError:
+        print("Warning: This version of tfds doesn't have the latest even_splits impl. "
+              "Please update or use tfds-nightly for better fine-grained split behaviour.")
+        has_buggy_even_splits = True
+except ImportError as e:
+    print(e)
+    print("Please install tensorflow_datasets package `pip install tensorflow-datasets`.")
+    exit(1)
+from .parser import Parser
+
+
+MAX_TP_SIZE = 8  # maximum TF threadpool size, only doing jpeg decodes and queuing activities
+SHUFFLE_SIZE = 8192  # examples to shuffle in DS queue
+PREFETCH_SIZE = 2048  # examples to prefetch
+
+
+def even_split_indices(split, n, num_examples):
+    partitions = [round(i * num_examples / n) for i in range(n + 1)]
+    return [f"{split}[{partitions[i]}:{partitions[i + 1]}]" for i in range(n)]
+
+
+def get_class_labels(info):
+    if 'label' not in info.features:
+        return {}
+    class_label = info.features['label']
+    class_to_idx = {n: class_label.str2int(n) for n in class_label.names}
+    return class_to_idx
+
+
+class ParserTfds(Parser):
+    """ Wrap Tensorflow Datasets for use in PyTorch
+
+    There several things to be aware of:
+      * To prevent excessive examples being dropped per epoch w/ distributed training or multiplicity of
+         dataloader workers, the train iterator wraps to avoid returning partial batches that trigger drop_last
+         https://github.com/pytorch/pytorch/issues/33413
+      * With PyTorch IterableDatasets, each worker in each replica operates in isolation, the final batch
+        from each worker could be a different size. For training this is worked around by option above, for
+        validation extra examples are inserted iff distributed mode is enabled so that the batches being reduced
+        across replicas are of same size. This will slightly alter the results, distributed validation will not be
+        100% correct. This is similar to common handling in DistributedSampler for normal Datasets but a bit worse
+        since there are up to N * J extra examples with IterableDatasets.
+      * The sharding (splitting of dataset into TFRecord) files imposes limitations on the number of
+        replicas and dataloader workers you can use. For really small datasets that only contain a few shards
+        you may have to train non-distributed w/ 1-2 dataloader workers. This is likely not a huge concern as the
+        benefit of distributed training or fast dataloading should be much less for small datasets.
+      * This wrapper is currently configured to return individual, decompressed image examples from the TFDS
+        dataset. The augmentation (transforms) and batching is still done in PyTorch. It would be possible
+        to specify TF augmentation fn and return augmented batches w/ some modifications to other downstream
+        components.
+
+    """
+
+    def __init__(
+            self,
+            root,
+            name,
+            split='train',
+            is_training=False,
+            batch_size=None,
+            download=False,
+            repeats=0,
+            seed=42,
+            input_name='image',
+            input_image='RGB',
+            target_name='label',
+            target_image='',
+            prefetch_size=None,
+            shuffle_size=None,
+            max_threadpool_size=None
+    ):
+        """ Tensorflow-datasets Wrapper
+
+        Args:
+            root: root data dir (ie your TFDS_DATA_DIR. not dataset specific sub-dir)
+            name: tfds dataset name (eg `imagenet2012`)
+            split: tfds dataset split (can use all TFDS split strings eg `train[:10%]`)
+            is_training: training mode, shuffle enabled, dataset len rounded by batch_size
+            batch_size: batch_size to use to unsure total examples % batch_size == 0 in training across all dis nodes
+            download: download and build TFDS dataset if set, otherwise must use tfds CLI
+            repeats: iterate through (repeat) the dataset this many times per iteration (once if 0 or 1)
+            seed: common seed for shard shuffle across all distributed/worker instances
+            input_name: name of Feature to return as data (input)
+            input_image: image mode if input is an image (currently PIL mode string)
+            target_name: name of Feature to return as target (label)
+            target_image: image mode if target is an image (currently PIL mode string)
+            prefetch_size: override default tf.data prefetch buffer size
+            shuffle_size: override default tf.data shuffle buffer size
+            max_threadpool_size: override default threadpool size for tf.data
+        """
+        super().__init__()
+        self.root = root
+        self.split = split
+        self.is_training = is_training
+        if self.is_training:
+            assert batch_size is not None, \
+                "Must specify batch_size in training mode for reasonable behaviour w/ TFDS wrapper"
+        self.batch_size = batch_size
+        self.repeats = repeats
+        self.common_seed = seed  # a seed that's fixed across all worker / distributed instances
+
+        # performance settings
+        self.prefetch_size = prefetch_size or PREFETCH_SIZE
+        self.shuffle_size = shuffle_size or SHUFFLE_SIZE
+        self.max_threadpool_size = max_threadpool_size or MAX_TP_SIZE
+
+        # TFDS builder and split information
+        self.input_name = input_name  # FIXME support tuples / lists of inputs and targets and full range of Feature
+        self.input_image = input_image
+        self.target_name = target_name
+        self.target_image = target_image
+        self.builder = tfds.builder(name, data_dir=root)
+        # NOTE: the tfds command line app can be used download & prepare datasets if you don't enable download flag
+        if download:
+            self.builder.download_and_prepare()
+        self.class_to_idx = get_class_labels(self.builder.info) if self.target_name == 'label' else {}
+        self.split_info = self.builder.info.splits[split]
+        self.num_examples = self.split_info.num_examples
+
+        # Distributed world state
+        self.dist_rank = 0
+        self.dist_num_replicas = 1
+        if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+            self.dist_rank = dist.get_rank()
+            self.dist_num_replicas = dist.get_world_size()
+
+        # Attributes that are updated in _lazy_init, including the tf.data pipeline itself
+        self.global_num_workers = 1
+        self.worker_info = None
+        self.worker_seed = 0  # seed unique to each work instance
+        self.subsplit = None  # set when data is distributed across workers using sub-splits
+        self.ds = None  # initialized lazily on each dataloader worker process
+
+    def _lazy_init(self):
+        """ Lazily initialize the dataset.
+
+        This is necessary to init the Tensorflow dataset pipeline in the (dataloader) process that
+        will be using the dataset instance. The __init__ method is called on the main process,
+        this will be called in a dataloader worker process.
+
+        NOTE: There will be problems if you try to re-use this dataset across different loader/worker
+        instances once it has been initialized. Do not call any dataset methods that can call _lazy_init
+        before it is passed to dataloader.
+        """
+        worker_info = torch.utils.data.get_worker_info()
+
+        # setup input context to split dataset across distributed processes
+        num_workers = 1
+        global_worker_id = 0
+        if worker_info is not None:
+            self.worker_info = worker_info
+            self.worker_seed = worker_info.seed
+            num_workers = worker_info.num_workers
+            self.global_num_workers = self.dist_num_replicas * num_workers
+            global_worker_id = self.dist_rank * num_workers + worker_info.id
+
+            """ Data sharding
+            InputContext will assign subset of underlying TFRecord files to each 'pipeline' if used.
+            My understanding is that using split, the underling TFRecord files will shuffle (shuffle_files=True)
+            between the splits each iteration, but that understanding could be wrong.
+
+            I am currently using a mix of InputContext shard assignment and fine-grained sub-splits for distributing
+            the data across workers. For training InputContext is used to assign shards to nodes unless num_shards
+            in dataset < total number of workers. Otherwise sub-split API is used for datasets without enough shards or
+            for validation where we can't drop examples and need to avoid minimize uneven splits to avoid padding.
+            """
+            should_subsplit = self.global_num_workers > 1 and (
+                    self.split_info.num_shards < self.global_num_workers or not self.is_training)
+            if should_subsplit:
+                # split the dataset w/o using sharding for more even examples / worker, can result in less optimal
+                # read patterns for distributed training (overlap across shards) so better to use InputContext there
+                if has_buggy_even_splits:
+                    # my even_split workaround doesn't work on subsplits, upgrade tfds!
+                    if not isinstance(self.split_info, tfds.core.splits.SubSplitInfo):
+                        subsplits = even_split_indices(self.split, self.global_num_workers, self.num_examples)
+                        self.subsplit = subsplits[global_worker_id]
+                else:
+                    subsplits = tfds.even_splits(self.split, self.global_num_workers)
+                    self.subsplit = subsplits[global_worker_id]
+
+        input_context = None
+        if self.global_num_workers > 1 and self.subsplit is None:
+            # set input context to divide shards among distributed replicas
+            input_context = tf.distribute.InputContext(
+                num_input_pipelines=self.global_num_workers,
+                input_pipeline_id=global_worker_id,
+                num_replicas_in_sync=self.dist_num_replicas  # FIXME does this arg have any impact?
+            )
+        read_config = tfds.ReadConfig(
+            shuffle_seed=self.common_seed,
+            shuffle_reshuffle_each_iteration=True,
+            input_context=input_context)
+        ds = self.builder.as_dataset(
+            split=self.subsplit or self.split, shuffle_files=self.is_training, read_config=read_config)
+        # avoid overloading threading w/ combo of TF ds threads + PyTorch workers
+        options = tf.data.Options()
+        thread_member = 'threading' if hasattr(options, 'threading') else 'experimental_threading'
+        getattr(options, thread_member).private_threadpool_size = max(1, self.max_threadpool_size // num_workers)
+        getattr(options, thread_member).max_intra_op_parallelism = 1
+        ds = ds.with_options(options)
+        if self.is_training or self.repeats > 1:
+            # to prevent excessive drop_last batch behaviour w/ IterableDatasets
+            # see warnings at https://pytorch.org/docs/stable/data.html#multi-process-data-loading
+            ds = ds.repeat()  # allow wrap around and break iteration manually
+        if self.is_training:
+            ds = ds.shuffle(min(self.num_examples, self.shuffle_size) // self.global_num_workers, seed=self.worker_seed)
+        ds = ds.prefetch(min(self.num_examples // self.global_num_workers, self.prefetch_size))
+        self.ds = tfds.as_numpy(ds)
+
+    def __iter__(self):
+        if self.ds is None:
+            self._lazy_init()
+
+        # Compute a rounded up sample count that is used to:
+        #   1. make batches even cross workers & replicas in distributed validation.
+        #     This adds extra examples and will slightly alter validation results.
+        #   2. determine loop ending condition in training w/ repeat enabled so that only full batch_size
+        #     batches are produced (underlying tfds iter wraps around)
+        target_example_count = math.ceil(max(1, self.repeats) * self.num_examples / self.global_num_workers)
+        if self.is_training:
+            # round up to nearest batch_size per worker-replica
+            target_example_count = math.ceil(target_example_count / self.batch_size) * self.batch_size
+
+        # Iterate until exhausted or sample count hits target when training (ds.repeat enabled)
+        example_count = 0
+        for example in self.ds:
+            input_data = example[self.input_name]
+            if self.input_image:
+                input_data = Image.fromarray(input_data, mode=self.input_image)
+            target_data = example[self.target_name]
+            if self.target_image:
+                target_data = Image.fromarray(target_data, mode=self.target_image)
+            yield input_data, target_data
+            example_count += 1
+            if self.is_training and example_count >= target_example_count:
+                # Need to break out of loop when repeat() is enabled for training w/ oversampling
+                # this results in extra examples per epoch but seems more desirable than dropping
+                # up to N*J batches per epoch (where N = num distributed processes, and J = num worker processes)
+                break
+
+        # Pad across distributed nodes (make counts equal by adding examples)
+        if not self.is_training and self.dist_num_replicas > 1 and self.subsplit is not None and \
+                0 < example_count < target_example_count:
+            # Validation batch padding only done for distributed training where results are reduced across nodes.
+            # For single process case, it won't matter if workers return different batch sizes.
+            # If using input_context or % based splits, sample count can vary significantly across workers and this
+            # approach should not be used (hence disabled if self.subsplit isn't set).
+            while example_count < target_example_count:
+                yield input_data, target_data  # yield prev sample again
+                example_count += 1
+
+    def __len__(self):
+        # this is just an estimate and does not factor in extra examples added to pad batches based on
+        # complete worker & replica info (not available until init in dataloader).
+        return math.ceil(max(1, self.repeats) * self.num_examples / self.dist_num_replicas)
+
+    def _filename(self, index, basename=False, absolute=False):
+        assert False, "Not supported"  # no random access to examples
+
+    def filenames(self, basename=False, absolute=False):
+        """ Return all filenames in dataset, overrides base"""
+        if self.ds is None:
+            self._lazy_init()
+        names = []
+        for sample in self.ds:
+            if len(names) > self.num_examples:
+                break  # safety for ds.repeat() case
+            if 'file_name' in sample:
+                name = sample['file_name']
+            elif 'filename' in sample:
+                name = sample['filename']
+            elif 'id' in sample:
+                name = sample['id']
+            else:
+                assert False, "No supported name field present"
+            names.append(name)
+        return names

openpoints/dataset/pcqm4m/__init__.py

@@ -0,0 +1 @@
+from .data import *

openpoints/dataset/pcqm4m/data.py

@@ -0,0 +1,62 @@
+import numpy as np
+import torch
+from ..dataset_base import DatasetBase
+from ..graph_dataset import GraphDataset
+from ..graph_dataset import SVDEncodingsGraphDataset
+from ..graph_dataset import StructuralDataset
+
+class PCQM4MDataset(DatasetBase):
+    def __init__(self, 
+                 dataset_path         ,
+                 dataset_name = 'PCQM4M',
+                 **kwargs
+                 ):
+        super().__init__(dataset_name = dataset_name,
+                         **kwargs)
+        self.dataset_path = dataset_path
+    
+    @property
+    def dataset(self):
+        try:
+            return self._dataset
+        except AttributeError:
+            from ogb.lsc import PCQM4MDataset
+            from ogb.utils import smiles2graph
+            self._smiles2graph = smiles2graph
+            self._dataset = PCQM4MDataset(root = self.dataset_path, only_smiles=True)
+            return self._dataset
+
+    @property
+    def record_tokens(self):
+        try:
+            return self._record_tokens
+        except AttributeError:
+            split = {'training':'train', 
+                     'validation':'valid', 
+                     'test':'test'}[self.split]
+            self._record_tokens = self.dataset.get_idx_split()[split]
+            return self._record_tokens
+    
+    def read_record(self, token):
+        smiles, target = self.dataset[token]
+        graph = self._smiles2graph(smiles)
+        graph['num_nodes'] = np.array(graph['num_nodes'], dtype=np.int16)
+        graph['edges'] = graph.pop('edge_index').T.astype(np.int16)
+        graph['edge_features'] = graph.pop('edge_feat').astype(np.int16)
+        graph['node_features'] = graph.pop('node_feat').astype(np.int16)
+        graph['target'] = np.array(target, np.float32)
+        return graph
+
+
+
+class PCQM4MGraphDataset(GraphDataset,PCQM4MDataset):
+    pass
+
+class PCQM4MSVDGraphDataset(SVDEncodingsGraphDataset,PCQM4MDataset):
+    pass
+
+class PCQM4MStructuralGraphDataset(StructuralDataset,PCQM4MGraphDataset):
+    pass
+
+class PCQM4MStructuralSVDGraphDataset(StructuralDataset,PCQM4MSVDGraphDataset):
+    pass

openpoints/dataset/pcqm4mv2/__init__.py

@@ -0,0 +1 @@
+from .data import *

openpoints/dataset/pcqm4mv2/data.py

@@ -0,0 +1,87 @@
+"""
+https://github.com/shamim-hussain/egt_pytorch
+"""
+
+import numpy as np
+from ..dataset_base import DatasetBase
+from ..graph_dataset import GraphDataset
+from ..graph_dataset import SVDEncodingsGraphDataset
+from ..graph_dataset import StructuralDataset
+from ..build import DATASETS
+
+
+@DATASETS.register_module()
+class PCQM4Mv2Dataset(DatasetBase):
+    def __init__(self,
+                 dataset_path,
+                 dataset_name='PCQM4MV2',
+                 **kwargs
+                 ):
+        super().__init__(dataset_name=dataset_name,
+                         **kwargs)
+        self.dataset_path = dataset_path
+
+    @property
+    def dataset(self):
+        try:
+            return self._dataset
+        except AttributeError:
+            from ogb.lsc import PCQM4Mv2Dataset
+            # smiles: simplified molecular-input line-entry system (SMILES). sequence representation of molecules
+            # The molecules are provided as SMILES strings (sequence representation of molecules)
+            # i = 1234
+            # print(dataset[i]) # ('CC(NCC[C@H]([C@@H]1CCC(=CC1)C)C)C', 6.811009678015001)
+
+            # The second option provides a molecular graph object constructed from the SMILES string.
+            # After preprocessing, the file size will be around 8GB.
+            # # get i-th molecule and its target value (nan for test data)
+            # i = 1234
+            # print(dataset[i]) # (graph_obj, 6.811009678015001)
+
+            # details: https://ogb.stanford.edu/docs/lsc/pcqm4mv2/
+            from ogb.utils import smiles2graph
+            self._smiles2graph = smiles2graph
+            self._dataset = PCQM4Mv2Dataset(root=self.dataset_path, only_smiles=True)
+            return self._dataset
+
+    @property
+    def record_tokens(self):
+        try:
+            return self._record_tokens
+        except AttributeError:
+            split = {'training': 'train',
+                     'validation': 'valid',
+                     'test': 'test-dev',
+                     'challenge': 'test-challenge'}[self.split]
+            self._record_tokens = self.dataset.get_idx_split()[split]
+            return self._record_tokens
+
+    def read_record(self, token):
+        smiles, target = self.dataset[token]
+        graph = self._smiles2graph(smiles)
+        graph['num_nodes'] = np.array(graph['num_nodes'], dtype=np.int16)
+        graph['edges'] = graph.pop('edge_index').T.astype(np.int16)
+        graph['edge_features'] = graph.pop('edge_feat').astype(np.int16)
+        graph['node_features'] = graph.pop('node_feat').astype(np.int16)
+        graph['target'] = np.array(target, np.float32)
+        return graph
+
+
+@DATASETS.register_module()
+class PCQM4Mv2GraphDataset(GraphDataset, PCQM4Mv2Dataset):
+    pass
+
+
+@DATASETS.register_module()
+class PCQM4Mv2SVDGraphDataset(SVDEncodingsGraphDataset, PCQM4Mv2Dataset):
+    pass
+
+
+@DATASETS.register_module()
+class PCQM4Mv2StructuralGraphDataset(StructuralDataset, PCQM4Mv2GraphDataset):
+    pass
+
+
+@DATASETS.register_module()
+class PCQM4Mv2StructuralSVDGraphDataset(StructuralDataset, PCQM4Mv2SVDGraphDataset):
+    pass

openpoints/dataset/s3dis/__init__.py

@@ -0,0 +1,2 @@
+from .s3dis_sphere import S3DISSphere 
+from .s3dis import S3DIS

openpoints/dataset/s3dis/s3dis.py

@@ -0,0 +1,156 @@
+import os
+import pickle
+import logging
+import numpy as np
+from tqdm import tqdm
+import torch
+from torch.utils.data import Dataset
+from ..data_util import crop_pc, voxelize
+from ..build import DATASETS
+
+
+@DATASETS.register_module()
+class S3DIS(Dataset):
+    classes = ['ceiling',
+               'floor',
+               'wall',
+               'beam',
+               'column',
+               'window',
+               'door',
+               'chair',
+               'table',
+               'bookcase',
+               'sofa',
+               'board',
+               'clutter']
+    num_classes = 13
+    num_per_class = np.array([3370714, 2856755, 4919229, 318158, 375640, 478001, 974733,
+                              650464, 791496, 88727, 1284130, 229758, 2272837], dtype=np.int32)
+    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],
+                   'table':       [170, 120, 200],
+                   'chair':       [255, 0, 0],
+                   'sofa':        [200, 100, 100],
+                   'bookcase':    [10, 200, 100],
+                   'board':       [200, 200, 200],
+                   'clutter':     [50, 50, 50]}
+    cmap = [*class2color.values()]
+    gravity_dim = 2
+    """S3DIS dataset, loading the subsampled entire room as input without block/sphere subsampling.
+    number of points per room in average, median, and std: (794855.5, 1005913.0147058824, 939501.4733064277)
+    Args:
+        data_root (str, optional): Defaults to 'data/S3DIS/s3disfull'.
+        test_area (int, optional): Defaults to 5.
+        voxel_size (float, optional): the voxel size for donwampling. Defaults to 0.04.
+        voxel_max (_type_, optional): subsample the max number of point per point cloud. Set None to use all points.  Defaults to None.
+        split (str, optional): Defaults to 'train'.
+        transform (_type_, optional): Defaults to None.
+        loop (int, optional): split loops for each epoch. Defaults to 1.
+        presample (bool, optional): wheter to downsample each point cloud before training. Set to False to downsample on-the-fly. Defaults to False.
+        variable (bool, optional): where to use the original number of points. The number of point per point cloud is variable. Defaults to False.
+    """
+    def __init__(self,
+                 data_root: str = 'data/S3DIS/s3disfull',
+                 test_area: int = 5,
+                 voxel_size: float = 0.04,
+                 voxel_max=None,
+                 split: str = 'train',
+                 transform=None,
+                 loop: int = 1,
+                 presample: bool = False,
+                 variable: bool = False,
+                 shuffle: bool = True,
+                 ):
+
+        super().__init__()
+        self.split, self.voxel_size, self.transform, self.voxel_max, self.loop = \
+            split, voxel_size, transform, voxel_max, loop
+        self.presample = presample
+        self.variable = variable
+        self.shuffle = shuffle
+
+        raw_root = os.path.join(data_root, 'raw')
+        self.raw_root = raw_root
+        data_list = sorted(os.listdir(raw_root))
+        data_list = [item[:-4] for item in data_list if 'Area_' in item]
+        if split == 'train':
+            self.data_list = [
+                item for item in data_list if not 'Area_{}'.format(test_area) in item]
+        else:
+            self.data_list = [
+                item for item in data_list if 'Area_{}'.format(test_area) in item]
+
+        processed_root = os.path.join(data_root, 'processed')
+        filename = os.path.join(
+            processed_root, f's3dis_{split}_area{test_area}_{voxel_size:.3f}_{str(voxel_max)}.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 S3DISFull {split} split on Test Area {test_area}'):
+                data_path = os.path.join(raw_root, item + '.npy')
+                cdata = np.load(data_path).astype(np.float32)
+                cdata[:, :3] -= np.min(cdata[:, :3], 0)
+                if voxel_size:
+                    coord, feat, label = cdata[:,0:3], cdata[:, 3:6], cdata[:, 6:7]
+                    uniq_idx = voxelize(coord, voxel_size)
+                    coord, feat, label = coord[uniq_idx], feat[uniq_idx], label[uniq_idx]
+                    cdata = np.hstack((coord, feat, label))
+                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")
+        self.data_idx = np.arange(len(self.data_list))
+        assert len(self.data_idx) > 0
+        logging.info(f"\nTotally {len(self.data_idx)} samples in {split} set")
+
+    def __getitem__(self, idx):
+        data_idx = self.data_idx[idx % len(self.data_idx)]
+        if self.presample:
+            coord, feat, label = np.split(self.data[data_idx], [3, 6], axis=1)
+        else:
+            data_path = os.path.join(
+                self.raw_root, self.data_list[data_idx] + '.npy')
+            cdata = np.load(data_path).astype(np.float32)
+            cdata[:, :3] -= np.min(cdata[:, :3], 0)
+            coord, feat, label = cdata[:, :3], cdata[:, 3:6], cdata[:, 6:7]
+            coord, feat, label = crop_pc(
+                coord, feat, label, self.split, self.voxel_size, self.voxel_max,
+                downsample=not self.presample, variable=self.variable, shuffle=self.shuffle)
+            # TODO: do we need to -np.min in cropped data?
+        label = label.squeeze(-1).astype(np.long)
+        data = {'pos': coord, 'x': feat, 'y': label}
+        # pre-process.
+        if self.transform is not None:
+            data = self.transform(data)
+
+        if 'heights' not in data.keys():
+            data['heights'] =  torch.from_numpy(coord[:, self.gravity_dim:self.gravity_dim+1].astype(np.float32))
+        return data
+
+    def __len__(self):
+        return len(self.data_idx) * self.loop
+        # return 1   # debug
+
+
+"""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)
+vis_multi_points([old_data['pos'][:, :3], data['pos'][:, :3].numpy()], colors=[old_data['x'][:, :3]/255.,data['x'][:, :3].numpy()])
+"""