openpoints 0.1.0__py3-none-any.whl

openpoints/dataset/data_util.py

@@ -0,0 +1,192 @@
+from curses import keyname
+import numpy as np
+import torch
+import os
+import os.path as osp
+import ssl
+import sys
+import urllib
+import h5py
+from typing import Optional
+
+
+class IO:
+    @classmethod
+    def get(cls, file_path):
+        _, file_extension = os.path.splitext(file_path)
+
+        if file_extension in ['.npy']:
+            return cls._read_npy(file_path)
+        elif file_extension in ['.pcd']:
+            return cls._read_pcd(file_path)
+        elif file_extension in ['.h5']:
+            return cls._read_h5(file_path)
+        elif file_extension in ['.txt']:
+            return cls._read_txt(file_path)
+        else:
+            raise Exception('Unsupported file extension: %s' % file_extension)
+
+    # References: https://github.com/numpy/numpy/blob/master/numpy/lib/format.py
+    @classmethod
+    def _read_npy(cls, file_path):
+        return np.load(file_path)
+    # # References: https://github.com/dimatura/pypcd/blob/master/pypcd/pypcd.py#L275
+    # # Support PCD files without compression ONLY!
+    # @classmethod
+    # def _read_pcd(cls, file_path):
+    #     pc = open3d.io.read_point_cloud(file_path)
+    #     ptcloud = np.array(pc.points)
+    #     return ptcloud
+
+    @classmethod
+    def _read_txt(cls, file_path):
+        return np.loadtxt(file_path)
+
+    @classmethod
+    def _read_h5(cls, file_path):
+        f = h5py.File(file_path, 'r')
+        return f['data'][()]
+
+
+# download
+def download_url(url: str, folder: str, log: bool = True,
+                 filename: Optional[str] = None):
+    r"""Downloads the content of an URL to a specific folder. 
+    Borrowed from https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/data/download.py
+    Args:
+        url (string): The url.
+        folder (string): The folder.
+        log (bool, optional): If :obj:`False`, will not print anything to the
+            console. (default: :obj:`True`)
+    """
+
+    if filename is None:
+        filename = url.rpartition('/')[2]
+        filename = filename if filename[0] == '?' else filename.split('?')[0]
+
+    path = osp.join(folder, filename)
+
+    if osp.exists(path):  # pragma: no cover
+        if log:
+            print(f'Using existing file {filename}', file=sys.stderr)
+        return path
+
+    if log:
+        print(f'Downloading {url}', file=sys.stderr)
+
+    os.makedirs(folder, exist_ok=True)
+    context = ssl._create_unverified_context()
+    data = urllib.request.urlopen(url, context=context)
+
+    with open(path, 'wb') as f:
+        # workaround for https://bugs.python.org/issue42853
+        while True:
+            chunk = data.read(10 * 1024 * 1024)
+            if not chunk:
+                break
+            f.write(chunk)
+
+    return path
+
+
+def fnv_hash_vec(arr):
+    """
+    FNV64-1A
+    """
+    assert arr.ndim == 2
+    arr = arr.copy()
+    arr -= arr.min(0)
+    arr = arr.astype(np.uint64, copy=False)
+    hashed_arr = np.uint64(14695981039346656037) * \
+        np.ones(arr.shape[0], dtype=np.uint64)
+    for j in range(arr.shape[1]):
+        hashed_arr *= np.uint64(1099511628211)
+        hashed_arr = np.bitwise_xor(hashed_arr, arr[:, j])
+    return hashed_arr
+
+
+def ravel_hash_vec(arr):
+    """
+    Ravel the coordinates after subtracting the min coordinates.
+    """
+    assert arr.ndim == 2
+    arr = arr.copy()
+    arr -= arr.min(0)
+    arr = arr.astype(np.uint64, copy=False)
+    arr_max = arr.max(0).astype(np.uint64) + 1
+
+    keys = np.zeros(arr.shape[0], dtype=np.uint64)
+    # Fortran style indexing
+    for j in range(arr.shape[1] - 1):
+        keys += arr[:, j]
+        keys *= arr_max[j + 1]
+    keys += arr[:, -1]
+    return keys
+
+
+def voxelize(coord, voxel_size=0.05, hash_type='fnv', mode=0):
+    discrete_coord = np.floor(coord / np.array(voxel_size))
+    if hash_type == 'ravel':
+        key = ravel_hash_vec(discrete_coord)
+    else:
+        key = fnv_hash_vec(discrete_coord)
+
+    idx_sort = np.argsort(key)
+    key_sort = key[idx_sort]
+    _, voxel_idx, count = np.unique(key_sort, return_counts=True, return_inverse=True)
+    if mode == 0:  # train mode
+        idx_select = np.cumsum(np.insert(count, 0, 0)[
+                               0:-1]) + np.random.randint(0, count.max(), count.size) % count
+        idx_unique = idx_sort[idx_select]
+        return idx_unique
+    else:  # val mode
+        return idx_sort, voxel_idx, count
+
+
+def crop_pc(coord, feat, label, split='train',
+            voxel_size=0.04, voxel_max=None,
+            downsample=True, variable=True, shuffle=True):
+    if voxel_size and downsample:
+        # Is this shifting a must? I borrow it from Stratified Transformer and Point Transformer. 
+        coord -= coord.min(0) 
+        uniq_idx = voxelize(coord, voxel_size)
+        coord, feat, label = coord[uniq_idx], feat[uniq_idx] if feat is not None else None, label[uniq_idx] if label is not None else None
+    if voxel_max is not None:
+        crop_idx = None
+        N = len(label)  # the number of points
+        if N >= voxel_max:
+            init_idx = np.random.randint(N) if 'train' in split else N // 2
+            crop_idx = np.argsort(
+                np.sum(np.square(coord - coord[init_idx]), 1))[:voxel_max]
+        elif not variable:
+            # fill more points for non-variable case (batched data)
+            cur_num_points = N
+            query_inds = np.arange(cur_num_points)
+            padding_choice = np.random.choice(
+                cur_num_points, voxel_max - cur_num_points)
+            crop_idx = np.hstack([query_inds, query_inds[padding_choice]])
+        crop_idx = np.arange(coord.shape[0]) if crop_idx is None else crop_idx
+        if shuffle:
+            shuffle_choice = np.random.permutation(np.arange(len(crop_idx)))
+            crop_idx = crop_idx[shuffle_choice]
+        coord, feat, label = coord[crop_idx], feat[crop_idx] if feat is not None else None, label[crop_idx] if label is not None else None
+    coord -= coord.min(0) 
+    return coord.astype(np.float32), feat.astype(np.float32) if feat is not None else None , label.astype(np.long) if label is not None else None
+
+
+def get_features_by_keys(data, keys='pos,x'):
+    key_list = keys.split(',')
+    if len(key_list) == 1:
+        return data[keys].transpose(1,2).contiguous()
+    else:
+        return torch.cat([data[key] for key in keys.split(',')], -1).transpose(1,2).contiguous()
+
+
+def get_class_weights(num_per_class, normalize=False):
+    weight = num_per_class / float(sum(num_per_class))
+    ce_label_weight = 1 / (weight + 0.02)
+
+    if normalize:
+        ce_label_weight = (ce_label_weight *
+                           len(ce_label_weight)) / ce_label_weight.sum()
+    return torch.from_numpy(ce_label_weight.astype(np.float32))

openpoints/dataset/datalist.py

@@ -0,0 +1,67 @@
+
+import torch
+from torch.utils.data import dataset
+from tqdm import tqdm
+from pathlib import Path
+import numpy as np
+
+
+class DataList(dataset.Dataset):
+    def __init__(self, 
+                 dataset_name,
+                 split, 
+                 data_list,  
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.dataset_name = dataset_name
+        self.data_list = data_list 
+        self.split = split
+        
+    def load_data(self, data_path):
+        if 's3dis' in self.dataset_name:
+            data = np.load(data_path)  # xyzrgbl, N*7
+            coord, feat, label = data[:, :3], data[:, 3:6], data[:, 6]
+            feat = np.clip(feat / 255., 0, 1).astype(np.float32)
+        elif 'scannet' in self.dataset_name:
+            data = torch.load(data_path)  # xyzrgbl, N*7
+            if self.split != 'test':
+                coord, feat, label = data[0], data[1], data[2]
+            else:
+                coord, feat, label = data[0], data[1], None
+            feat = np.clip((feat + 1) / 2., 0, 1).astype(np.float32)
+        elif 'semantickitti' in self.dataset_name: 
+            points = self.load_pc_kitti(pc_path)
+            labels = self.load_label_kitti(label_path, self.remap_lut_read)
+            
+        coord -= coord.min(0)
+
+        idx_points = []
+        voxel_size = cfg.dataset.common.get('voxel_size', None)
+        if voxel_size is not None:
+            idx_sort, count = voxelize(coord, voxel_size, mode=1)
+            for i in range(count.max()):
+                idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + i % count
+                idx_part = idx_sort[idx_select]
+                np.random.shuffle(idx_part)
+                idx_points.append(idx_part)
+        else:
+            idx_points.append(np.arange(label.shape[0]))
+        return coord, feat, label, idx_points
+
+
+    
+    def __len__(self):
+        return len(self.record_tokens)
+    
+    def __getitem__(self, index):
+        token  = self.record_tokens[index]
+        try:
+            return self._records[token]
+        except AttributeError:
+            record = self.read_record(token)
+            self._records = {token:record}
+            return record
+        except KeyError:
+            record = self.read_record(token)
+            self._records[token] = record
+            return record

openpoints/dataset/dataset_base.py

@@ -0,0 +1,96 @@
+
+import torch
+from torch.utils.data import dataset
+from tqdm import tqdm
+from pathlib import Path
+
+
+class DatasetBase(dataset.Dataset):
+    def __init__(self, 
+                 dataset_name, split,
+                 cache_dir = None, 
+                 load_cache_if_exists=True, 
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.dataset_name = dataset_name
+        self.split = split
+        self.cache_dir = cache_dir
+        
+        self.is_cached = False
+        if load_cache_if_exists:
+            self.cache(verbose=0, must_exist=True)
+        
+    @property
+    def record_tokens(self):
+        raise NotImplementedError
+    
+    def read_record(self, token):
+        raise NotImplementedError
+    
+    def __len__(self):
+        return len(self.record_tokens)
+    
+    def __getitem__(self, index):
+        token  = self.record_tokens[index]
+        try:
+            return self._records[token]
+        except AttributeError:
+            record = self.read_record(token)
+            self._records = {token:record}
+            return record
+        except KeyError:
+            record = self.read_record(token)
+            self._records[token] = record
+            return record
+    
+    def read_all_records(self, verbose=1):
+        self._records = {}
+        if verbose:
+            print(f'Reading all {self.split} records...', flush=True)
+            for token in tqdm(self.record_tokens):
+                self._records[token] = self.read_record(token)
+        else:
+            for token in self.record_tokens:
+                self._records[token] = self.read_record(token)
+    
+    def get_cache_path(self, path=None):
+        if path is None: path = self.cache_dir
+        base_path = (Path(path)/self.dataset_name)/self.split
+        base_path.mkdir(parents=True, exist_ok=True)
+        return base_path
+    
+    def cache_load_and_save(self, base_path, op, verbose):
+        tokens_path = base_path/'tokens.pt'
+        records_path = base_path/'records.pt'
+        
+        if op == 'load':
+            self._record_tokens = torch.load(str(tokens_path))
+            self._records = torch.load(str(records_path))
+        elif op == 'save':
+            if tokens_path.exists() and records_path.exists() \
+                and hasattr(self, '_record_tokens') and hasattr(self, '_records'):
+                return
+            self.read_all_records(verbose=verbose)
+            torch.save(self.record_tokens, str(tokens_path))
+            torch.save(self._records, str(records_path))
+        else:
+            raise ValueError(f'Unknown operation: {op}')
+    
+    def cache(self, path=None, verbose=1, must_exist=False):
+        if self.is_cached: return
+        
+        base_path = self.get_cache_path(path)
+        try:
+            if verbose: print(f'Trying to load {self.split} cache from disk...', flush=True)
+            self.cache_load_and_save(base_path, 'load', verbose)
+            if verbose: print(f'Loaded {self.split} cache from disk.', flush=True)
+        except FileNotFoundError:
+            if must_exist: return
+            
+            if verbose: print(f'{self.split} cache does not exist! Cacheing...', flush=True)
+            self.cache_load_and_save(base_path, 'save', verbose)
+            if verbose: print(f'Saved {self.split} cache to disk.', flush=True)
+        
+        self.is_cached = True
+
+

openpoints/dataset/graph_dataset/__init__.py

@@ -0,0 +1,3 @@
+from .graph_dataset import GraphDataset, graphdata_collate
+from .svd_encodings_dataset import SVDEncodingsGraphDataset
+from .structural_dataset import StructuralDataset

openpoints/dataset/graph_dataset/graph_dataset.py

@@ -0,0 +1,93 @@
+
+import torch
+import numpy as np
+
+from ..dataset_base import DatasetBase
+
+from .stack_with_pad import stack_with_pad
+from collections import defaultdict
+from numba.typed import List
+
+
+class GraphDataset(DatasetBase):
+    def __init__(self,
+                 num_nodes_key     = 'num_nodes',
+                 edges_key         = 'edges',
+                 node_features_key = 'node_features',
+                 edge_features_key = 'edge_features',
+                 node_mask_key     = 'node_mask',
+                 targets_key       = 'target',
+                 include_node_mask = True,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.num_nodes_key     = num_nodes_key    
+        self.edges_key         = edges_key        
+        self.node_features_key = node_features_key
+        self.edge_features_key = edge_features_key
+        self.node_mask_key     = node_mask_key    
+        self.targets_key       = targets_key      
+        self.include_node_mask = include_node_mask
+    
+    def __getitem__(self, index):
+        item = super().__getitem__(index)
+        if self.include_node_mask:
+            item = item.copy()
+            item[self.node_mask_key] = np.ones((item[self.num_nodes_key],), dtype=np.uint8)
+        return item
+    
+    def _calculate_max_nodes(self):
+        self._max_nodes = self[0][self.num_nodes_key]
+        self._max_nodes_index = 0
+        for i in range(1, super().__len__()):
+            graph = super().__getitem__(i)
+            cur_nodes = graph[self.num_nodes_key]
+            if cur_nodes > self._max_nodes:
+                self._max_nodes = cur_nodes
+                self._max_nodes_index = i
+    
+    @property
+    def max_nodes(self):
+        try:
+            return self._max_nodes
+        except AttributeError:
+            self._calculate_max_nodes()
+            return self._max_nodes
+    
+    @property
+    def max_nodes_index(self):
+        try:
+            return self._max_nodes_index
+        except AttributeError:
+            self._calculate_max_nodes()
+            return self._max_nodes_index
+    
+    def cache_load_and_save(self, base_path, op, verbose):
+        super().cache_load_and_save(base_path, op, verbose)
+        max_nodes_path = base_path/'max_nodes_data.pt'
+        
+        if op == 'load':
+            max_nodes_data = torch.load(str(max_nodes_path))
+            self._max_nodes = max_nodes_data['max_nodes']
+            self._max_nodes_index = max_nodes_data['max_nodes_index']
+        elif op == 'save':
+            if verbose: print(f'Calculating {self.split} max nodes...',flush=True)
+            max_nodes_data = {'max_nodes': self.max_nodes,
+                              'max_nodes_index': self.max_nodes_index}
+            torch.save(max_nodes_data, str(max_nodes_path))
+        else:
+            raise ValueError(f'Unknown operation: {op}')
+    
+    def max_batch(self, batch_size, collate_fn):
+        return collate_fn([self.__getitem__(self.max_nodes_index)] * batch_size)
+
+
+
+def graphdata_collate(batch):
+    batch_data = defaultdict(List)
+    for elem in batch:
+        for k,v in elem.items():
+            batch_data[k].append(v)
+    
+    out = {k:torch.from_numpy(stack_with_pad(dat)) 
+                    for k, dat in batch_data.items()}
+    return out

openpoints/dataset/graph_dataset/stack_with_pad.py

@@ -0,0 +1,91 @@
+import numpy as np
+import numba as nb
+
+
+@nb.njit
+def stack_with_pad_4d(inputs):
+    num_elem = len(inputs)
+    ms_0, ms_1, ms_2, ms_3 = inputs[0].shape
+    
+    for i in range(1,num_elem):
+        is_0, is_1, is_2, is_3 = inputs[i].shape
+        ms_0 = max(is_0, ms_0)
+        ms_1 = max(is_1, ms_1)
+        ms_2 = max(is_2, ms_2)
+        ms_3 = max(is_3, ms_3)
+    
+    stacked_shape = (num_elem,ms_0,ms_1,ms_2,ms_3)
+    stacked = np.zeros(stacked_shape, dtype=inputs[0].dtype)
+    
+    for i, elem in enumerate(inputs):
+        stacked[i][:elem.shape[0],:elem.shape[1],:elem.shape[2],:elem.shape[3]] = elem
+    return stacked
+
+@nb.njit
+def stack_with_pad_3d(inputs):
+    num_elem = len(inputs)
+    ms_0, ms_1, ms_2 = inputs[0].shape
+    
+    for i in range(1,num_elem):
+        is_0, is_1, is_2 = inputs[i].shape
+        ms_0 = max(is_0, ms_0)
+        ms_1 = max(is_1, ms_1)
+        ms_2 = max(is_2, ms_2)
+    
+    stacked_shape = (num_elem,ms_0,ms_1,ms_2)
+    stacked = np.zeros(stacked_shape, dtype=inputs[0].dtype)
+    
+    for i, elem in enumerate(inputs):
+        stacked[i][:elem.shape[0],:elem.shape[1],:elem.shape[2]] = elem
+    return stacked
+
+@nb.njit
+def stack_with_pad_2d(inputs):
+    num_elem = len(inputs)
+    ms_0, ms_1 = inputs[0].shape
+    
+    for i in range(1,num_elem):
+        is_0, is_1 = inputs[i].shape
+        ms_0 = max(is_0, ms_0)
+        ms_1 = max(is_1, ms_1)
+    
+    stacked_shape = (num_elem,ms_0,ms_1)
+    stacked = np.zeros(stacked_shape, dtype=inputs[0].dtype)
+    
+    for i, elem in enumerate(inputs):
+        stacked[i][:elem.shape[0],:elem.shape[1]] = elem
+    return stacked
+
+@nb.njit
+def stack_with_pad_1d(inputs):
+    num_elem = len(inputs)
+    ms_0 = inputs[0].shape[0]
+    
+    for i in range(1,num_elem):
+        is_0 = inputs[i].shape[0]
+        ms_0 = max(is_0, ms_0)
+    
+    stacked_shape = (num_elem,ms_0)
+    stacked = np.zeros(stacked_shape, dtype=inputs[0].dtype)
+    
+    for i, elem in enumerate(inputs):
+        stacked[i][:elem.shape[0]] = elem
+    return stacked
+
+
+def stack_with_pad(inputs):
+    shape_rank = np.ndim(inputs[0])
+    if shape_rank == 0:
+        return np.stack(inputs)
+    if shape_rank == 1:
+        return stack_with_pad_1d(inputs)
+    elif shape_rank == 2:
+        return stack_with_pad_2d(inputs)
+    elif shape_rank == 3:
+        return stack_with_pad_3d(inputs)
+    elif shape_rank == 4:
+        return stack_with_pad_4d(inputs)
+    else:
+        raise ValueError('Only support up to 4D tensor')
+
+

openpoints/dataset/graph_dataset/structural_dataset.py

@@ -0,0 +1,73 @@
+import numpy as np
+import numba as nb
+
+from .graph_dataset import GraphDataset
+
+NODE_FEATURES_OFFSET = 128
+EDGE_FEATURES_OFFSET = 8
+
+@nb.njit
+def floyd_warshall(A):
+    n = A.shape[0]
+    D = np.zeros((n,n), dtype=np.int16)
+    
+    for i in range(n):
+        for j in range(n):
+            if i == j:
+                pass
+            elif A[i,j] == 0:
+                D[i,j] = 510
+            else:
+                D[i,j] = 1
+    
+    for k in range(n):
+        for i in range(n):
+            for j in range(n):
+                old_dist = D[i,j]
+                new_dist = D[i,k] + D[k,j]
+                if new_dist < old_dist:
+                    D[i,j] = new_dist
+    return D
+
+@nb.njit
+def preprocess_data(num_nodes, edges, node_feats, edge_feats):
+    node_feats = node_feats + np.arange(1,node_feats.shape[-1]*NODE_FEATURES_OFFSET+1,
+                                            NODE_FEATURES_OFFSET,dtype=np.int16)
+    edge_feats = edge_feats + np.arange(1,edge_feats.shape[-1]*EDGE_FEATURES_OFFSET+1,
+                                            EDGE_FEATURES_OFFSET,dtype=np.int16)
+    
+    A = np.zeros((num_nodes,num_nodes),dtype=np.int16)
+    E = np.zeros((num_nodes,num_nodes,edge_feats.shape[-1]),dtype=np.int16)
+    for k in range(edges.shape[0]):
+        i,j = edges[k,0], edges[k,1]
+        A[i,j] = 1
+        E[i,j] = edge_feats[k]
+    
+    D = floyd_warshall(A)
+    return node_feats, D, E
+
+
+class StructuralDataset(GraphDataset):
+    def __init__(self,
+                 distance_matrix_key      = 'distance_matrix',
+                 feature_matrix_key       = 'feature_matrix',
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.distance_matrix_key      = distance_matrix_key      
+        self.feature_matrix_key       = feature_matrix_key  
+        
+    def __getitem__(self, index):
+        item = super().__getitem__(index)
+        
+        num_nodes = int(item[self.num_nodes_key])
+        edges = item.pop(self.edges_key)
+        node_feats = item.pop(self.node_features_key)
+        edge_feats = item.pop(self.edge_features_key)
+        
+        node_feats, dist_mat, edge_feats_mat = preprocess_data(num_nodes, edges, node_feats, edge_feats)
+        item[self.node_features_key] = node_feats
+        item[self.distance_matrix_key] = dist_mat
+        item[self.feature_matrix_key] = edge_feats_mat
+        
+        return item
+