openpoints 0.1.0__py3-none-any.whl

openpoints/dataset/matterport3d/matterport3d.py

@@ -0,0 +1,210 @@
+"""_summary_
+"""
+
+
+import os
+import numpy as np
+import pickle
+import logging
+import pathlib
+import glob
+import random
+import h5py
+import pandas as pd
+import torch
+from tqdm import tqdm
+from torch.utils.data import Dataset
+from ..build import DATASETS
+from openpoints.models.layers import fps
+import warnings
+warnings.filterwarnings('ignore')
+
+
+def pc_normalize(pc):
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
+    pc = pc / m
+    return pc
+
+
+@DATASETS.register_module()
+class MP40(Dataset):
+    classes = ["wall",
+               "floor",
+               "chair",
+               "door"
+               "table",
+               "picture",
+               "cabinet",
+               "cushion",
+               "window",
+               "sofa"
+               "bed"
+               "curtain",
+               "chest_of_drawers",
+               "plant",
+               "sink"
+               "stairs",
+               "ceiling",
+               "toilet",
+               "stool",
+               "towel",
+               "mirror",
+               "tv_monitor",
+               "shower",
+               "column",
+               "bathtub",
+               "counter",
+               "fireplace",
+               "lighting",
+               "beam"
+               "railing",
+               "shelving",
+               "blinds",
+               "gym_equipment",
+               "seating",
+               "board_panel",
+               "furniture",
+               "appliances",
+               "clothes",
+               "objects",
+               "misc"
+               ]
+
+    def __init__(self,
+                 data_dir,
+                 num_points=1024,
+                 split='train',
+                 transform=None,
+                 use_normal=False,
+                 **kwargs
+                 ):
+        self.npoints = num_points
+        self.preprocess = True
+        self.uniform = True
+        self.split = split
+        self.use_normal = use_normal
+        self.transform = transform
+
+        root_dir = pathlib.Path(data_dir)
+        data_dir = root_dir / 'raw'
+        data_list = root_dir / f'mattportobject_{split}_list.txt'
+
+        if not data_list.exists():
+            all_files = [i.name for i in data_dir.glob('*.npy')]
+            random.shuffle(all_files)
+            n_files = len(all_files)
+            n_train = int(0.8 * n_files)
+            n_val = int(0.1 * n_files)
+            train_files = all_files[:n_train]
+            val_files = all_files[n_train:n_train+n_val]
+            test_files = all_files[n_train+n_val:]
+            train_list = str(root_dir / (f'mattportobject_train_list.txt'))
+            val_list = str(root_dir / (f'mattportobject_val_list.txt'))
+            test_list = str(root_dir / (f'mattportobject_test_list.txt'))
+            with open(train_list, 'w') as f:
+                f.writelines('\n'.join(train_files))
+            with open(val_list, 'w') as f:
+                f.writelines('\n'.join(val_files))
+            with open(test_list, 'w') as f:
+                f.writelines('\n'.join(test_files))
+        with open(data_list, 'r') as f:
+            lines = f.read().splitlines()
+        self.datapath = [data_dir / l for l in lines]
+
+        if self.uniform:
+            self.save_path = os.path.join(
+                root_dir, 'matterport3dobjects_%s_%dpts_fps.h5' % (split, 2048))
+        else:
+            self.save_path = os.path.join(
+                root_dir, 'matterport3dobjects_%s_%dpts.h5' % (split, 2048))
+
+        npoints = 2048
+        if not os.path.exists(self.save_path):
+            logging.info(
+                'Processing data %s (only running in the first time)...' % self.save_path)
+            list_of_points = []
+            list_of_labels = []
+
+            for index in tqdm(range(len(self.datapath)), total=len(self.datapath)):
+                fn = self.datapath[index]
+                data = np.load(fn, allow_pickle=True).item()
+                cls = data['label']
+                point_set = data['points'].astype(np.float32)
+                if self.uniform:
+                    point_set = fps(torch.from_numpy(point_set).unsqueeze(
+                        0).cuda(), npoints).squeeze(0).cpu().numpy()
+                else:
+                    point_set = point_set[:npoints, :]
+
+                list_of_points.append(point_set)
+                list_of_labels.append(cls)
+
+            all_points = np.stack(list_of_points)
+            all_labels = np.stack(list_of_labels)
+            all_data = {'data': all_points, 'label': all_labels}
+            hf = h5py.File(self.save_path, 'w')
+            data = hf.create_group('data')
+            for k, v in all_data.items():
+                data[k] = v
+            hf.close()
+        logging.info('Load processed data from %s...' % self.save_path)
+        with h5py.File(self.save_path, 'r') as f:
+            data = f['data']['data'][:].astype('float32')
+            label = f['data']['label'][:].astype('int32')
+
+        # TODO: make this in preprocessing step. not here.
+        cls_mapping = pd.read_csv(str(pathlib.Path(__file__).parent / "category_mapping.tsv"),
+                                  skiprows=1, header=None, sep='\t', usecols=[0, 16]).values.astype(int)
+
+        # step 1 remove negative label (useless)
+        idx = np.argwhere(label > 0).squeeze()
+        data = data[idx]
+        label = label[idx]
+
+        # step 2 mapping label to mat40
+        # cls_mapping is a np array. row 0 = wall = label(wall) --> mpcat40
+        label = cls_mapping[label-1][:, 1]
+
+        # step 3, remove the label not in mat40 (0, and 41)
+        label_bigger_0 = label > 0
+        label_smaller_41 = label < 41
+        idx = []
+        for i in range(len(label)):
+            if label_bigger_0[i] and label_smaller_41[i]:
+                idx.append(i)
+        self.data = data[idx]
+        self.label = label[idx] - 1
+
+    def __len__(self):
+        return len(self.data)
+
+    @property
+    def num_classes(self):
+        return self.label.max() + 1
+
+    def __getitem__(self, index):
+        points, label = self.data[index][:self.npoints], self.label[index]
+        if self.split == 'train':
+            np.random.shuffle(points)
+        data = {'pos': points[:, :3],
+                'x': points[:, 3:6 + 3 * self.use_normal],
+                'y': label
+                }
+        if self.transform is not None:
+            data = self.transform(data)
+        """_summary_
+        from openpoints.dataset.vis3d import vis_points
+
+        for i, idx in enumerate(np.where(self.label==22)[0]):
+            if i >5:
+                break
+            vis_points(self.data[idx, :, :3], self.data[idx, :, 3:6]/255.)
+        """
+        if 'heights' in data.keys():
+            data['x'] = torch.cat(
+                (data['pos'], data['heights'], data['x']), dim=1)
+        else:
+            data['x'] = torch.cat((data['pos'], data['x']), dim=1)
+        return data

openpoints/dataset/matterport3d/matterport3d_dataprocessing.py

@@ -0,0 +1,105 @@
+
+import os, sys, zipfile, h5py, numpy as np, multiprocessing
+from pathlib import Path
+from tqdm import tqdm 
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))
+from openpoints.dataset.utils3d import *
+from openpoints.dataset.vis3d import *
+
+
+def unzip_folders(data_dir, unzipped_path):
+    for dirpath, dirs, files in os.walk(data_dir):	 
+        for file in files:
+            if file.endswith('zip'):
+                file_path = os.path.join(dirpath, file)
+                zip_file = zipfile.ZipFile(file_path)
+                for names in zip_file.namelist():
+                    zip_file.extract(names, unzipped_path)
+                zip_file.close() 
+
+
+def save_room_objects(mesh_path, save_name):
+    """save the objects in a rooom together 
+
+    Args:
+        mesh_path (_type_): _description_
+        save_name (_type_): _description_
+    """
+    all_points = []
+    all_ins_labels = []
+    all_cls_labels = []
+    
+    points, faces, f_cls_id, f_ins_id = read_mesh_vf(mesh_path)
+    # category id: https://github.com/niessner/Matterport/blob/master/metadata/category_mapping.tsv 
+    # vis_points(points[:, :3], points[:, 3:6])
+    f_cls_id_unique = np.unique(f_cls_id)   # the unique category ids (eg. pillow, wall)
+    for cls in f_cls_id_unique:
+        face_cls_flag = f_cls_id==cls   # wheter face is current cls
+        face_cls = faces[face_cls_flag] # pick the face which is current cls
+        face_cls_ins = f_ins_id[face_cls_flag]  # pick the instance label of the face
+        for ins_id in np.unique(face_cls_ins):  # for each instance 
+            point_idx = face_cls[face_cls_ins==ins_id].reshape(-1)
+            points_cls_ins = points[point_idx]
+            all_points.append(points_cls_ins)
+            all_ins_labels.append(ins_id)
+            all_cls_labels.append(cls)
+    data = {'points': all_points, 'ins_labels': all_ins_labels, 'y': all_cls_labels}
+    np.save(f'{save_name}.npy', data)
+
+
+def save_per_object(mesh_path, save_name):
+    """save each object in a h5 file.  
+
+    Args:
+        mesh_path (_type_): _description_
+        save_name (_type_): _description_
+    """
+    
+    points, faces, f_cls_id, f_ins_id = read_mesh_vf(mesh_path)
+    # category id: https://github.com/niessner/Matterport/blob/master/metadata/category_mapping.tsv 
+    # vis_points(points[:, :3], points[:, 3:6])
+    f_cls_id_unique = np.unique(f_cls_id)   # the unique category ids (eg. pillow, wall)
+    for cls in f_cls_id_unique:
+        face_cls_flag = f_cls_id==cls   # wheter face is current cls
+        face_cls = faces[face_cls_flag] # pick the face which is current cls
+        face_cls_ins = f_ins_id[face_cls_flag]  # pick the instance label of the face
+
+        for ins_id in np.unique(face_cls_ins):  # for each instance 
+            point_idx = face_cls[face_cls_ins==ins_id].reshape(-1)
+            points_cls_ins = points[point_idx]
+            vis_points(points_cls_ins[:, :3], points_cls_ins[:, 3:6])
+            data = {'points': points_cls_ins, 'ins_label': ins_id, 'label': cls}
+            
+            # TODO: debug comment for now. 
+            np.save(f'{save_name}_C{cls}_I{ins_id}.npy', data)
+    
+    
+if __name__ == "__main__":
+    # unzip data at first. 
+    # root_path = '/data/3D/Matterport3d/v1/scans'
+    # unzipped_path = '/data/3D/Matterport3d/v1/scans_unzipped'
+    # unzip_folders(root_path, unzipped_path)
+    
+    # loading all points to a single h5 file
+    
+    root_path = Path('/data/3D/Matterport3d/v1/')
+    
+    raw_path = root_path.joinpath('scans_unzipped')
+    save_dir = root_path.joinpath('objects')
+    save_dir.mkdir(exist_ok=True)
+    
+    all_points = []
+    all_ins_labels = []
+    all_cls_labels = []
+    
+    # p = multiprocessing.Pool()
+    for path in Path(raw_path).rglob('*.ply'):
+        file_path = path.parent.joinpath(path.name)
+        file_name = '/'.join(path.parts[-3:])
+        save_name = '_'.join([path.parts[-3], path.parts[-1].split('.')[0]])
+        save_path = os.path.join(save_dir, save_name)
+        print(f'===> processing {file_name} ...')
+        save_per_object(file_path, save_path)    
+    #     p.apply_async(save_per_object, [file_path, save_path]) 
+    # p.close()
+    # p.join() # Wait for all child processes to close.

openpoints/dataset/modelnet/__init__.py

@@ -0,0 +1,3 @@
+from .modelnet40_normal_resampled_loader import ModelNet
+from .modelnet40_ply_2048_loader import ModelNet40Ply2048
+

openpoints/dataset/modelnet/modelnet40_normal_resampled_loader.py

@@ -0,0 +1,124 @@
+'''
+Borrowed from PointBERT 
+@file: ModelNet.py
+@time: 2021/3/19 15:51
+'''
+import os
+import numpy as np
+import warnings
+import pickle
+from tqdm import tqdm
+from torch.utils.data import Dataset
+from ..build import DATASETS
+import torch
+import logging
+warnings.filterwarnings('ignore')
+
+
+def pc_normalize(pc):
+    centroid = np.mean(pc, axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
+    pc = pc / m
+    return pc
+
+
+def farthest_point_sample(point, npoint):
+    """
+    Input:
+        xyz: pointcloud data, [N, D]
+        npoint: number of samples
+    Return:
+        centroids: sampled pointcloud index, [npoint, D]
+    """
+    N, D = point.shape
+    xyz = point[:, :3]
+    centroids = np.zeros((npoint,))
+    distance = np.ones((N,)) * 1e10
+    farthest = np.random.randint(0, N)
+    for i in range(npoint):
+        centroids[i] = farthest
+        centroid = xyz[farthest, :]
+        dist = np.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = np.argmax(distance, -1)
+    point = point[centroids.astype(np.int32)]
+    return point
+
+
+@DATASETS.register_module()
+class ModelNet(Dataset):
+    def __init__(self,
+                 data_dir, num_points, num_classes,
+                 use_normals=False,
+                 split='train',
+                 transform=None
+                 ):
+        self.root = os.path.join(data_dir, 'modelnet40_normal_resampled')
+        self.npoints = num_points
+        self.use_normals = use_normals
+        self.num_category = num_classes
+        split = 'train' if split.lower() == 'train' else 'test'  # val = test
+        self.split = split
+
+        if self.num_category == 10:
+            self.catfile = os.path.join(
+                self.root, 'modelnet10_shape_names.txt')
+        else:
+            self.catfile = os.path.join(
+                self.root, 'modelnet40_shape_names.txt')
+
+        self.cat = [line.rstrip() for line in open(self.catfile)]
+        self.classes = dict(zip(self.cat, range(len(self.cat))))
+
+        shape_ids = {}
+        if self.num_category == 10:
+            shape_ids['train'] = [line.rstrip() for line in open(
+                os.path.join(self.root, 'modelnet10_train.txt'))]
+            shape_ids['test'] = [line.rstrip() for line in open(
+                os.path.join(self.root, 'modelnet10_test.txt'))]
+        else:
+            shape_ids['train'] = [line.rstrip() for line in open(
+                os.path.join(self.root, 'modelnet40_train.txt'))]
+            shape_ids['test'] = [line.rstrip() for line in open(
+                os.path.join(self.root, 'modelnet40_test.txt'))]
+
+        assert (split == 'train' or split == 'test')
+        shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]]
+        self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], shape_ids[split][i]) + '.txt') for i
+                         in range(len(shape_ids[split]))]
+        logging.info('The size of %s data is %d' % (split, len(self.datapath)))
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.datapath)
+
+    @property
+    def num_classes(self):
+        return self.list_of_labels.max() + 1
+
+    def _get_item(self, index):
+        fn = self.datapath[index]
+        cls = self.classes[self.datapath[index][0]]
+        label = np.array([cls]).astype(np.int64)
+        point_set = np.loadtxt(fn[1], delimiter=',').astype(np.float32)
+        return point_set[:, 0:3], point_set[:, 3:], label[0]
+
+    def __getitem__(self, index):
+        points, feats, label = self._get_item(index)
+        if self.split == 'train':
+            np.random.shuffle(points)
+        data = {'pos': points,
+                'y': label
+                }
+        if self.use_normals:
+            data['x'] = feats
+        if self.transform is not None:
+            data = self.transform(data)
+
+        if self.use_normals:
+            data['x'] = torch.cat((data['pos'], data['x']), dim=1)
+        if 'heights' in data.keys():
+            data['x'] = torch.cat((data['x'], data['heights']), dim=1)
+        return data

openpoints/dataset/modelnet/modelnet40_ply_2048_loader.py

@@ -0,0 +1,160 @@
+"""Modified from DeepGCN and DGCNN
+Reference: https://github.com/lightaime/deep_gcns_torch/tree/master/examples/classification
+"""
+import os
+import glob
+import h5py
+import numpy as np
+import pickle
+import logging
+import ssl
+import urllib
+from pathlib import Path
+from tqdm import tqdm
+import torch
+from torch.utils.data import Dataset
+from torchvision.datasets.utils import extract_archive, check_integrity
+from ..build import DATASETS
+
+
+def download_and_extract_archive(url, path, md5=None):
+    # Works when the SSL certificate is expired for the link
+    path = Path(path)
+    extract_path = path
+    if not path.exists():
+        path.mkdir(parents=True, exist_ok=True)
+        file_path = path / Path(url).name
+        if not file_path.exists() or not check_integrity(file_path, md5):
+            print(f'{file_path} not found or corrupted')
+            print(f'downloading from {url}')
+            context = ssl.SSLContext()
+            with urllib.request.urlopen(url, context=context) as response:
+                with tqdm(total=response.length) as pbar:
+                    with open(file_path, 'wb') as file:
+                        chunk_size = 1024
+                        chunks = iter(lambda: response.read(chunk_size), '')
+                        for chunk in chunks:
+                            if not chunk:
+                                break
+                            pbar.update(chunk_size)
+                            file.write(chunk)
+            extract_archive(str(file_path), str(extract_path))
+    return extract_path
+
+
+def load_data(data_dir, partition, url):
+    download_and_extract_archive(url, data_dir)
+    all_data = []
+    all_label = []
+    for h5_name in glob.glob(os.path.join(data_dir, 'modelnet40_ply_hdf5_2048', 'ply_data_%s*.h5' % partition)):
+        with h5py.File(h5_name, 'r') as f:
+            data = f['data'][:].astype('float32')
+            label = f['label'][:].astype('int64')
+        all_data.append(data)
+        all_label.append(label)
+    all_data = np.concatenate(all_data, axis=0)
+    all_label = np.concatenate(all_label, axis=0).squeeze(-1)
+    return all_data, all_label
+
+
+@DATASETS.register_module()
+class ModelNet40Ply2048(Dataset):
+    """
+    This is the data loader for ModelNet 40
+    ModelNet40 contains 12,311 meshed CAD models from 40 categories.
+    num_points: 1024 by default
+    data_dir
+    paritition: train or test
+    """
+    dir_name = 'modelnet40_ply_hdf5_2048'
+    md5 = 'c9ab8e6dfb16f67afdab25e155c79e59'
+    url = f'https://shapenet.cs.stanford.edu/media/{dir_name}.zip'
+    classes = ['airplane',
+               'bathtub',
+               'bed',
+               'bench',
+               'bookshelf',
+               'bottle',
+               'bowl',
+               'car',
+               'chair',
+               'cone',
+               'cup',
+               'curtain',
+               'desk',
+               'door',
+               'dresser',
+               'flower_pot',
+               'glass_box',
+               'guitar',
+               'keyboard',
+               'lamp',
+               'laptop',
+               'mantel',
+               'monitor',
+               'night_stand',
+               'person',
+               'piano',
+               'plant',
+               'radio',
+               'range_hood',
+               'sink',
+               'sofa',
+               'stairs',
+               'stool',
+               'table',
+               'tent',
+               'toilet',
+               'tv_stand',
+               'vase',
+               'wardrobe',
+               'xbox']
+
+    def __init__(self,
+                 num_points=1024,
+                 data_dir="./data/ModelNet40Ply2048",
+                 split='train',
+                 transform=None
+                 ):
+        data_dir = os.path.join(
+            os.getcwd(), data_dir) if data_dir.startswith('.') else data_dir
+        self.partition = 'train' if split.lower() == 'train' else 'test'  # val = test
+        self.data, self.label = load_data(data_dir, self.partition, self.url)
+        self.num_points = num_points
+        logging.info(f'==> sucessfully loaded {self.partition} data')
+        self.transform = transform
+
+    def __getitem__(self, item):
+        pointcloud = self.data[item][:self.num_points]
+        label = self.label[item]
+
+        if self.partition == 'train':
+            np.random.shuffle(pointcloud)
+        data = {'pos': pointcloud,
+                'y': label
+                }
+        if self.transform is not None:
+            data = self.transform(data)
+
+        if 'heights' in data.keys():
+            data['x'] = torch.cat((data['pos'], data['heights']), dim=1)
+        else:
+            data['x'] = data['pos']
+        return data
+
+    def __len__(self):
+        return self.data.shape[0]
+
+    @property
+    def num_classes(self):
+        return np.max(self.label) + 1
+
+    """ for visulalization
+    from openpoints.dataset import vis_multi_points
+    import copy
+    old_points = copy.deepcopy(data['pos'])
+    if self.transform is not None:
+        data = self.transform(data)
+    new_points = copy.deepcopy(data['pos'])
+    vis_multi_points([old_points, new_points.numpy()])
+    End of visulization """

openpoints/dataset/molhiv/__init__.py

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

openpoints/dataset/molhiv/data.py

@@ -0,0 +1,59 @@
+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 MOLHIVDataset(DatasetBase):
+    def __init__(self, 
+                 dataset_path             ,
+                 dataset_name = 'MOLHIV' ,
+                 **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.graphproppred import GraphPropPredDataset
+            self._dataset = GraphPropPredDataset(name='ogbg-molhiv', root=self.dataset_path)
+            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):
+        graph, target = self.dataset[token]
+        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 MOLHIVGraphDataset(GraphDataset,MOLHIVDataset):
+    pass
+
+class MOLHIVSVDGraphDataset(SVDEncodingsGraphDataset,MOLHIVDataset):
+    pass
+
+class MOLHIVStructuralGraphDataset(StructuralDataset,MOLHIVGraphDataset):
+    pass
+
+class MOLHIVStructuralSVDGraphDataset(StructuralDataset,MOLHIVSVDGraphDataset):
+    pass

openpoints/dataset/molpcba/__init__.py

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