Repository: rhett-chen/graspness_implementation Branch: main Commit: ff33da111e72 Files: 53 Total size: 198.5 KB Directory structure: gitextract_d4wahnnt/ ├── .gitignore ├── LICENSE ├── README.md ├── command_test.sh ├── command_train.sh ├── dataset/ │ ├── generate_graspness.py │ ├── graspnet_dataset.py │ ├── simplify_dataset.py │ └── vis_graspness.py ├── doc/ │ └── example_data/ │ └── meta.mat ├── infer_vis_grasp.py ├── knn/ │ ├── knn_modules.py │ ├── setup.py │ └── src/ │ ├── cpu/ │ │ ├── knn_cpu.cpp │ │ └── vision.h │ ├── cuda/ │ │ ├── knn.cu │ │ └── vision.h │ ├── knn.h │ └── vision.cpp ├── models/ │ ├── backbone_resunet14.py │ ├── graspnet.py │ ├── loss.py │ ├── modules.py │ └── resnet.py ├── pointnet2/ │ ├── _ext_src/ │ │ ├── include/ │ │ │ ├── ball_query.h │ │ │ ├── cuda_utils.h │ │ │ ├── cylinder_query.h │ │ │ ├── group_points.h │ │ │ ├── interpolate.h │ │ │ ├── sampling.h │ │ │ └── utils.h │ │ └── src/ │ │ ├── ball_query.cpp │ │ ├── ball_query_gpu.cu │ │ ├── bindings.cpp │ │ ├── cylinder_query.cpp │ │ ├── cylinder_query_gpu.cu │ │ ├── group_points.cpp │ │ ├── group_points_gpu.cu │ │ ├── interpolate.cpp │ │ ├── interpolate_gpu.cu │ │ ├── sampling.cpp │ │ └── sampling_gpu.cu │ ├── pointnet2_modules.py │ ├── pointnet2_utils.py │ ├── pytorch_utils.py │ └── setup.py ├── requirements.txt ├── test.py ├── train.py └── utils/ ├── collision_detector.py ├── data_utils.py ├── label_generation.py └── loss_utils.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ **/__pycache__/** *.ipynb **/.ipynb_checkpoints/** *.npy *.npz **/.vscode/** **/grasp_label*/** **/log*/** **/dump*/** **/build/** *.o *.so *.egg **/*.egg-info/** logs dataset/tolerance **/.idea/ ================================================ FILE: LICENSE ================================================ GRASPNET-BASELINE SOFTWARE LICENSE AGREEMENT ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR DOWNLOAD THE SOFTWARE. This is a license agreement ("Agreement") between your academic institution or non-profit organization or self (called "Licensee" or "You" in this Agreement) and Shanghai Jiao Tong University (called "Licensor" in this Agreement). All rights not specifically granted to you in this Agreement are reserved for Licensor. RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive ownership of any copy of the Software (as defined below) licensed under this Agreement and hereby grants to Licensee a personal, non-exclusive, non-transferable license to use the Software for noncommercial research purposes, without the right to sublicense, pursuant to the terms and conditions of this Agreement. As used in this Agreement, the term "Software" means (i) the actual copy of all or any portion of code for program routines made accessible to Licensee by Licensor pursuant to this Agreement, inclusive of backups, updates, and/or merged copies permitted hereunder or subsequently supplied by Licensor, including all or any file structures, programming instructions, user interfaces and screen formats and sequences as well as any and all documentation and instructions related to it, and (ii) all or any derivatives and/or modifications created or made by You to any of the items specified in (i). CONFIDENTIALITY: Licensee acknowledges that the Software is proprietary to Licensor, and as such, Licensee agrees to receive all such materials in confidence and use the Software only in accordance with the terms of this Agreement. Licensee agrees to use reasonable effort to protect the Software from unauthorized use, reproduction, distribution, or publication. PERMITTED USES: The Software may be used for your own noncommercial internal research purposes. You understand and agree that Licensor is not obligated to implement any suggestions and/or feedback you might provide regarding the Software, but to the extent Licensor does so, you are not entitled to any compensation related thereto. DERIVATIVES: You may create derivatives of or make modifications to the Software, however, You agree that all and any such derivatives and modifications will be owned by Licensor and become a part of the Software licensed to You under this Agreement. You may only use such derivatives and modifications for your own noncommercial internal research purposes, and you may not otherwise use, distribute or copy such derivatives and modifications in violation of this Agreement. BACKUPS: If Licensee is an organization, it may make that number of copies of the Software necessary for internal noncommercial use at a single site within its organization provided that all information appearing in or on the original labels, including the copyright and trademark notices are copied onto the labels of the copies. USES NOT PERMITTED: You may not distribute, copy or use the Software except as explicitly permitted herein. Licensee has not been granted any trademark license as part of this Agreement and may not use the name or mark “AlphaPose", "Shanghai Jiao Tong" or any renditions thereof without the prior written permission of Licensor. You may not sell, rent, lease, sublicense, lend, time-share or transfer, in whole or in part, or provide third parties access to prior or present versions (or any parts thereof) of the Software. ASSIGNMENT: You may not assign this Agreement or your rights hereunder without the prior written consent of Licensor. Any attempted assignment without such consent shall be null and void. TERM: The term of the license granted by this Agreement is from Licensee's acceptance of this Agreement by downloading the Software or by using the Software until terminated as provided below. The Agreement automatically terminates without notice if you fail to comply with any provision of this Agreement. Licensee may terminate this Agreement by ceasing using the Software. Upon any termination of this Agreement, Licensee will delete any and all copies of the Software. You agree that all provisions which operate to protect the proprietary rights of Licensor shall remain in force should breach occur and that the obligation of confidentiality described in this Agreement is binding in perpetuity and, as such, survives the term of the Agreement. FEE: Provided Licensee abides completely by the terms and conditions of this Agreement, there is no fee due to Licensor for Licensee's use of the Software in accordance with this Agreement. DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON-INFRINGEMENT. LICENSEE BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND RELATED MATERIALS. SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is provided as part of this Agreement. EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent permitted under applicable law, Licensor shall not be liable for direct, indirect, special, incidental, or consequential damages or lost profits related to Licensee's use of and/or inability to use the Software, even if Licensor is advised of the possibility of such damage. EXPORT REGULATION: Licensee agrees to comply with any and all applicable U.S. export control laws, regulations, and/or other laws related to embargoes and sanction programs administered by the Office of Foreign Assets Control. SEVERABILITY: If any provision(s) of this Agreement shall be held to be invalid, illegal, or unenforceable by a court or other tribunal of competent jurisdiction, the validity, legality and enforceability of the remaining provisions shall not in any way be affected or impaired thereby. NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right or remedy under this Agreement shall be construed as a waiver of any future or other exercise of such right or remedy by Licensor. ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and entire agreement between Licensee and Licensor as to the matter set forth herein and supersedes any previous agreements, understandings, and arrangements between the parties relating hereto. ************************************************************************ THIRD-PARTY SOFTWARE NOTICES AND INFORMATION This project incorporates material from the project(s) listed below (collectively, "Third Party Code"). This Third Party Code is licensed to you under their original license terms set forth below. We reserves all other rights not expressly granted, whether by implication, estoppel or otherwise. 1. PyTorch (https://github.com/pytorch/pytorch) THIRD-PARTY SOFTWARE NOTICES AND INFORMATION This project incorporates material from the project(s) listed below (collectively, "Third Party Code"). This Third Party Code is licensed to you under their original license terms set forth below. We reserves all other rights not expressly granted, whether by implication, estoppel or otherwise. From PyTorch: Copyright (c) 2016- Facebook, Inc (Adam Paszke) Copyright (c) 2014- Facebook, Inc (Soumith Chintala) Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) Copyright (c) 2011-2013 NYU (Clement Farabet) Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) Copyright (c) 2006 Idiap Research Institute (Samy Bengio) Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) From Caffe2: Copyright (c) 2016-present, Facebook Inc. All rights reserved. All contributions by Facebook: Copyright (c) 2016 Facebook Inc. All contributions by Google: Copyright (c) 2015 Google Inc. All rights reserved. All contributions by Yangqing Jia: Copyright (c) 2015 Yangqing Jia All rights reserved. All contributions by Kakao Brain: Copyright 2019-2020 Kakao Brain All contributions from Caffe: Copyright(c) 2013, 2014, 2015, the respective contributors All rights reserved. All other contributions: Copyright(c) 2015, 2016 the respective contributors All rights reserved. Caffe2 uses a copyright model similar to Caffe: each contributor holds copyright over their contributions to Caffe2. The project versioning records all such contribution and copyright details. If a contributor wants to further mark their specific copyright on a particular contribution, they should indicate their copyright solely in the commit message of the change when it is committed. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America and IDIAP Research Institute nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 2. VoteNet (https://github.com/facebookresearch/votenet) MIT License Copyright (c) Facebook, Inc. and its affiliates. Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ************END OF THIRD-PARTY SOFTWARE NOTICES AND INFORMATION********** ================================================ FILE: README.md ================================================ # GraspNet graspness My implementation of paper "Graspness Discovery in Clutters for Fast and Accurate Grasp Detection" (ICCV 2021). [[paper](https://openaccess.thecvf.com/content/ICCV2021/papers/Wang_Graspness_Discovery_in_Clutters_for_Fast_and_Accurate_Grasp_Detection_ICCV_2021_paper.pdf)] [[dataset](https://graspnet.net/)] [[API](https://github.com/graspnet/graspnetAPI)] ## Requirements - Python 3 - PyTorch 1.8 - Open3d 0.8 - TensorBoard 2.3 - NumPy - SciPy - Pillow - tqdm - MinkowskiEngine ## Installation Get the code. ```bash git clone https://github.com/rhett-chen/graspness_implementation.git cd graspnet-graspness ``` Install packages via Pip. ```bash pip install -r requirements.txt ``` Compile and install pointnet2 operators (code adapted from [votenet](https://github.com/facebookresearch/votenet)). ```bash cd pointnet2 python setup.py install ``` Compile and install knn operator (code adapted from [pytorch_knn_cuda](https://github.com/chrischoy/pytorch_knn_cuda)). ```bash cd knn python setup.py install ``` Install graspnetAPI for evaluation. ```bash git clone https://github.com/graspnet/graspnetAPI.git cd graspnetAPI pip install . ``` For MinkowskiEngine, please refer https://github.com/NVIDIA/MinkowskiEngine ## Point level Graspness Generation Point level graspness label are not included in the original dataset, and need additional generation. Make sure you have downloaded the orginal dataset from [GraspNet](https://graspnet.net/). The generation code is in [dataset/generate_graspness.py](dataset/generate_graspness.py). ```bash cd dataset python generate_graspness.py --dataset_root /data3/graspnet --camera_type kinect ``` ## Simplify dataset original dataset grasp_label files have redundant data, We can significantly save the memory cost. The code is in [dataset/simplify_dataset.py](dataset/simplify_dataset.py) ```bash cd dataset python simplify_dataset.py --dataset_root /data3/graspnet ``` ## Training and Testing Training examples are shown in [command_train.sh](command_train.sh). `--dataset_root`, `--camera` and `--log_dir` should be specified according to your settings. You can use TensorBoard to visualize training process. Testing examples are shown in [command_test.sh](command_test.sh), which contains inference and result evaluation. `--dataset_root`, `--camera`, `--checkpoint_path` and `--dump_dir` should be specified according to your settings. Set `--collision_thresh` to -1 for fast inference. ## Results Results "In repo" report the model performance of my results without collision detection. Evaluation results on Kinect camera: | | | Seen | | | Similar | | | Novel | | |:--------:|:------:|:----------------:|:----------------:|:------:|:----------------:|:----------------:|:------:|:----------------:|:----------------:| | | __AP__ | AP0.8 | AP0.4 | __AP__ | AP0.8 | AP0.4 | __AP__ | AP0.8 | AP0.4 | | In paper | 61.19 | 71.46 | 56.04 | 47.39 | 56.78 | 40.43 | 19.01 | 23.73 | 10.60 | | In repo | 61.83 | 73.28 | 54.14 | 51.13 | 62.53 | 41.57 | 19.94 | 24.90 | 11.02 | ## Troubleshooting If you meet the torch.floor error in MinkowskiEngine, you can simply solve it by changing the source code of MinkowskiEngine: MinkowskiEngine/utils/quantization.py 262，from discrete_coordinates =_auto_floor(coordinates) to discrete_coordinates = coordinates ## Acknowledgement My code is mainly based on Graspnet-baseline https://github.com/graspnet/graspnet-baseline. ================================================ FILE: command_test.sh ================================================ CUDA_VISIBLE_DEVICES=4 python test.py --camera kinect --dump_dir logs/log_kn/dump_epoch10 --checkpoint_path logs/log_kn/minkresunet_epoch10.tar --batch_size 1 --dataset_root /data3/graspnet --infer --eval --collision_thresh -1 ================================================ FILE: command_train.sh ================================================ CUDA_VISIBLE_DEVICES=4 python train.py --camera kinect --log_dir logs/log_kn --batch_size 4 --learning_rate 0.001 --model_name minkuresunet --dataset_root /data3/graspnet ================================================ FILE: dataset/generate_graspness.py ================================================ import numpy as np import os from PIL import Image import scipy.io as scio import sys ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(ROOT_DIR) from utils.data_utils import get_workspace_mask, CameraInfo, create_point_cloud_from_depth_image from knn.knn_modules import knn import torch from graspnetAPI.utils.xmlhandler import xmlReader from graspnetAPI.utils.utils import get_obj_pose_list, transform_points import argparse parser = argparse.ArgumentParser() parser.add_argument('--dataset_root', default=None, required=True) parser.add_argument('--camera_type', default='kinect', help='Camera split [realsense/kinect]') if __name__ == '__main__': cfgs = parser.parse_args() dataset_root = cfgs.dataset_root # set dataset root camera_type = cfgs.camera_type # kinect / realsense save_path_root = os.path.join(dataset_root, 'graspness') num_views, num_angles, num_depths = 300, 12, 4 fric_coef_thresh = 0.8 point_grasp_num = num_views * num_angles * num_depths for scene_id in range(100): save_path = os.path.join(save_path_root, 'scene_' + str(scene_id).zfill(4), camera_type) if not os.path.exists(save_path): os.makedirs(save_path) labels = np.load( os.path.join(dataset_root, 'collision_label', 'scene_' + str(scene_id).zfill(4), 'collision_labels.npz')) collision_dump = [] for j in range(len(labels)): collision_dump.append(labels['arr_{}'.format(j)]) for ann_id in range(256): # get scene point cloud print('generating scene: {} ann: {}'.format(scene_id, ann_id)) depth = np.array(Image.open(os.path.join(dataset_root, 'scenes', 'scene_' + str(scene_id).zfill(4), camera_type, 'depth', str(ann_id).zfill(4) + '.png'))) seg = np.array(Image.open(os.path.join(dataset_root, 'scenes', 'scene_' + str(scene_id).zfill(4), camera_type, 'label', str(ann_id).zfill(4) + '.png'))) meta = scio.loadmat(os.path.join(dataset_root, 'scenes', 'scene_' + str(scene_id).zfill(4), camera_type, 'meta', str(ann_id).zfill(4) + '.mat')) intrinsic = meta['intrinsic_matrix'] factor_depth = meta['factor_depth'] camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2], factor_depth) cloud = create_point_cloud_from_depth_image(depth, camera, organized=True) # remove outlier and get objectness label depth_mask = (depth > 0) camera_poses = np.load(os.path.join(dataset_root, 'scenes', 'scene_' + str(scene_id).zfill(4), camera_type, 'camera_poses.npy')) camera_pose = camera_poses[ann_id] align_mat = np.load(os.path.join(dataset_root, 'scenes', 'scene_' + str(scene_id).zfill(4), camera_type, 'cam0_wrt_table.npy')) trans = np.dot(align_mat, camera_pose) workspace_mask = get_workspace_mask(cloud, seg, trans=trans, organized=True, outlier=0.02) mask = (depth_mask & workspace_mask) cloud_masked = cloud[mask] objectness_label = seg[mask] # get scene object and grasp info scene_reader = xmlReader(os.path.join(dataset_root, 'scenes', 'scene_' + str(scene_id).zfill(4), camera_type, 'annotations', '%04d.xml' % ann_id)) pose_vectors = scene_reader.getposevectorlist() obj_list, pose_list = get_obj_pose_list(camera_pose, pose_vectors) grasp_labels = {} for i in obj_list: file = np.load(os.path.join(dataset_root, 'grasp_label', '{}_labels.npz'.format(str(i).zfill(3)))) grasp_labels[i] = (file['points'].astype(np.float32), file['offsets'].astype(np.float32), file['scores'].astype(np.float32)) grasp_points = [] grasp_points_graspness = [] for i, (obj_idx, trans_) in enumerate(zip(obj_list, pose_list)): sampled_points, offsets, fric_coefs = grasp_labels[obj_idx] collision = collision_dump[i] # Npoints * num_views * num_angles * num_depths num_points = sampled_points.shape[0] valid_grasp_mask = ((fric_coefs <= fric_coef_thresh) & (fric_coefs > 0) & ~collision) valid_grasp_mask = valid_grasp_mask.reshape(num_points, -1) graspness = np.sum(valid_grasp_mask, axis=1) / point_grasp_num target_points = transform_points(sampled_points, trans_) target_points = transform_points(target_points, np.linalg.inv(camera_pose)) # fix bug grasp_points.append(target_points) grasp_points_graspness.append(graspness.reshape(num_points, 1)) grasp_points = np.vstack(grasp_points) grasp_points_graspness = np.vstack(grasp_points_graspness) grasp_points = torch.from_numpy(grasp_points).cuda() grasp_points_graspness = torch.from_numpy(grasp_points_graspness).cuda() grasp_points = grasp_points.transpose(0, 1).contiguous().unsqueeze(0) masked_points_num = cloud_masked.shape[0] cloud_masked_graspness = np.zeros((masked_points_num, 1)) part_num = int(masked_points_num / 10000) for i in range(1, part_num + 2): # lack of cuda memory if i == part_num + 1: cloud_masked_partial = cloud_masked[10000 * part_num:] if len(cloud_masked_partial) == 0: break else: cloud_masked_partial = cloud_masked[10000 * (i - 1):(i * 10000)] cloud_masked_partial = torch.from_numpy(cloud_masked_partial).cuda() cloud_masked_partial = cloud_masked_partial.transpose(0, 1).contiguous().unsqueeze(0) nn_inds = knn(grasp_points, cloud_masked_partial, k=1).squeeze() - 1 cloud_masked_graspness[10000 * (i - 1):(i * 10000)] = torch.index_select( grasp_points_graspness, 0, nn_inds).cpu().numpy() max_graspness = np.max(cloud_masked_graspness) min_graspness = np.min(cloud_masked_graspness) cloud_masked_graspness = (cloud_masked_graspness - min_graspness) / (max_graspness - min_graspness) np.save(os.path.join(save_path, str(ann_id).zfill(4) + '.npy'), cloud_masked_graspness) ================================================ FILE: dataset/graspnet_dataset.py ================================================ """ GraspNet dataset processing. Author: chenxi-wang """ import os import numpy as np import scipy.io as scio from PIL import Image import torch import collections.abc as container_abcs from torch.utils.data import Dataset from tqdm import tqdm import MinkowskiEngine as ME from data_utils import CameraInfo, transform_point_cloud, create_point_cloud_from_depth_image, get_workspace_mask class GraspNetDataset(Dataset): def __init__(self, root, grasp_labels=None, camera='kinect', split='train', num_points=20000, voxel_size=0.005, remove_outlier=True, augment=False, load_label=True): assert (num_points <= 50000) self.root = root self.split = split self.voxel_size = voxel_size self.num_points = num_points self.remove_outlier = remove_outlier self.grasp_labels = grasp_labels self.camera = camera self.augment = augment self.load_label = load_label self.collision_labels = {} if split == 'train': self.sceneIds = list(range(100)) elif split == 'test': self.sceneIds = list(range(100, 190)) elif split == 'test_seen': self.sceneIds = list(range(100, 130)) elif split == 'test_similar': self.sceneIds = list(range(130, 160)) elif split == 'test_novel': self.sceneIds = list(range(160, 190)) self.sceneIds = ['scene_{}'.format(str(x).zfill(4)) for x in self.sceneIds] self.depthpath = [] self.labelpath = [] self.metapath = [] self.scenename = [] self.frameid = [] self.graspnesspath = [] for x in tqdm(self.sceneIds, desc='Loading data path and collision labels...'): for img_num in range(256): self.depthpath.append(os.path.join(root, 'scenes', x, camera, 'depth', str(img_num).zfill(4) + '.png')) self.labelpath.append(os.path.join(root, 'scenes', x, camera, 'label', str(img_num).zfill(4) + '.png')) self.metapath.append(os.path.join(root, 'scenes', x, camera, 'meta', str(img_num).zfill(4) + '.mat')) self.graspnesspath.append(os.path.join(root, 'graspness', x, camera, str(img_num).zfill(4) + '.npy')) self.scenename.append(x.strip()) self.frameid.append(img_num) if self.load_label: collision_labels = np.load(os.path.join(root, 'collision_label', x.strip(), 'collision_labels.npz')) self.collision_labels[x.strip()] = {} for i in range(len(collision_labels)): self.collision_labels[x.strip()][i] = collision_labels['arr_{}'.format(i)] def scene_list(self): return self.scenename def __len__(self): return len(self.depthpath) def augment_data(self, point_clouds, object_poses_list): # Flipping along the YZ plane if np.random.random() > 0.5: flip_mat = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]]) point_clouds = transform_point_cloud(point_clouds, flip_mat, '3x3') for i in range(len(object_poses_list)): object_poses_list[i] = np.dot(flip_mat, object_poses_list[i]).astype(np.float32) # Rotation along up-axis/Z-axis rot_angle = (np.random.random() * np.pi / 3) - np.pi / 6 # -30 ~ +30 degree c, s = np.cos(rot_angle), np.sin(rot_angle) rot_mat = np.array([[1, 0, 0], [0, c, -s], [0, s, c]]) point_clouds = transform_point_cloud(point_clouds, rot_mat, '3x3') for i in range(len(object_poses_list)): object_poses_list[i] = np.dot(rot_mat, object_poses_list[i]).astype(np.float32) return point_clouds, object_poses_list def __getitem__(self, index): if self.load_label: return self.get_data_label(index) else: return self.get_data(index) def get_data(self, index, return_raw_cloud=False): depth = np.array(Image.open(self.depthpath[index])) seg = np.array(Image.open(self.labelpath[index])) meta = scio.loadmat(self.metapath[index]) scene = self.scenename[index] try: intrinsic = meta['intrinsic_matrix'] factor_depth = meta['factor_depth'] except Exception as e: print(repr(e)) print(scene) camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2], factor_depth) # generate cloud cloud = create_point_cloud_from_depth_image(depth, camera, organized=True) # get valid points depth_mask = (depth > 0) if self.remove_outlier: camera_poses = np.load(os.path.join(self.root, 'scenes', scene, self.camera, 'camera_poses.npy')) align_mat = np.load(os.path.join(self.root, 'scenes', scene, self.camera, 'cam0_wrt_table.npy')) trans = np.dot(align_mat, camera_poses[self.frameid[index]]) workspace_mask = get_workspace_mask(cloud, seg, trans=trans, organized=True, outlier=0.02) mask = (depth_mask & workspace_mask) else: mask = depth_mask cloud_masked = cloud[mask] if return_raw_cloud: return cloud_masked # sample points random if len(cloud_masked) >= self.num_points: idxs = np.random.choice(len(cloud_masked), self.num_points, replace=False) else: idxs1 = np.arange(len(cloud_masked)) idxs2 = np.random.choice(len(cloud_masked), self.num_points - len(cloud_masked), replace=True) idxs = np.concatenate([idxs1, idxs2], axis=0) cloud_sampled = cloud_masked[idxs] ret_dict = {'point_clouds': cloud_sampled.astype(np.float32), 'coors': cloud_sampled.astype(np.float32) / self.voxel_size, 'feats': np.ones_like(cloud_sampled).astype(np.float32), } return ret_dict def get_data_label(self, index): depth = np.array(Image.open(self.depthpath[index])) seg = np.array(Image.open(self.labelpath[index])) meta = scio.loadmat(self.metapath[index]) graspness = np.load(self.graspnesspath[index]) # for each point in workspace masked point cloud scene = self.scenename[index] try: obj_idxs = meta['cls_indexes'].flatten().astype(np.int32) poses = meta['poses'] intrinsic = meta['intrinsic_matrix'] factor_depth = meta['factor_depth'] except Exception as e: print(repr(e)) print(scene) camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2], factor_depth) # generate cloud cloud = create_point_cloud_from_depth_image(depth, camera, organized=True) # get valid points depth_mask = (depth > 0) if self.remove_outlier: camera_poses = np.load(os.path.join(self.root, 'scenes', scene, self.camera, 'camera_poses.npy')) align_mat = np.load(os.path.join(self.root, 'scenes', scene, self.camera, 'cam0_wrt_table.npy')) trans = np.dot(align_mat, camera_poses[self.frameid[index]]) workspace_mask = get_workspace_mask(cloud, seg, trans=trans, organized=True, outlier=0.02) mask = (depth_mask & workspace_mask) else: mask = depth_mask cloud_masked = cloud[mask] seg_masked = seg[mask] # sample points if len(cloud_masked) >= self.num_points: idxs = np.random.choice(len(cloud_masked), self.num_points, replace=False) else: idxs1 = np.arange(len(cloud_masked)) idxs2 = np.random.choice(len(cloud_masked), self.num_points - len(cloud_masked), replace=True) idxs = np.concatenate([idxs1, idxs2], axis=0) cloud_sampled = cloud_masked[idxs] seg_sampled = seg_masked[idxs] graspness_sampled = graspness[idxs] objectness_label = seg_sampled.copy() objectness_label[objectness_label > 1] = 1 object_poses_list = [] grasp_points_list = [] grasp_widths_list = [] grasp_scores_list = [] for i, obj_idx in enumerate(obj_idxs): if (seg_sampled == obj_idx).sum() < 50: continue object_poses_list.append(poses[:, :, i]) points, widths, scores = self.grasp_labels[obj_idx] collision = self.collision_labels[scene][i] # (Np, V, A, D) idxs = np.random.choice(len(points), min(max(int(len(points) / 4), 300), len(points)), replace=False) grasp_points_list.append(points[idxs]) grasp_widths_list.append(widths[idxs]) collision = collision[idxs].copy() scores = scores[idxs].copy() scores[collision] = 0 grasp_scores_list.append(scores) if self.augment: cloud_sampled, object_poses_list = self.augment_data(cloud_sampled, object_poses_list) ret_dict = {'point_clouds': cloud_sampled.astype(np.float32), 'coors': cloud_sampled.astype(np.float32) / self.voxel_size, 'feats': np.ones_like(cloud_sampled).astype(np.float32), 'graspness_label': graspness_sampled.astype(np.float32), 'objectness_label': objectness_label.astype(np.int64), 'object_poses_list': object_poses_list, 'grasp_points_list': grasp_points_list, 'grasp_widths_list': grasp_widths_list, 'grasp_scores_list': grasp_scores_list} return ret_dict def load_grasp_labels(root): obj_names = list(range(1, 89)) grasp_labels = {} for obj_name in tqdm(obj_names, desc='Loading grasping labels...'): label = np.load(os.path.join(root, 'grasp_label_simplified', '{}_labels.npz'.format(str(obj_name - 1).zfill(3)))) grasp_labels[obj_name] = (label['points'].astype(np.float32), label['width'].astype(np.float32), label['scores'].astype(np.float32)) return grasp_labels def minkowski_collate_fn(list_data): coordinates_batch, features_batch = ME.utils.sparse_collate([d["coors"] for d in list_data], [d["feats"] for d in list_data]) coordinates_batch, features_batch, _, quantize2original = ME.utils.sparse_quantize( coordinates_batch, features_batch, return_index=True, return_inverse=True) res = { "coors": coordinates_batch, "feats": features_batch, "quantize2original": quantize2original } def collate_fn_(batch): if type(batch[0]).__module__ == 'numpy': return torch.stack([torch.from_numpy(b) for b in batch], 0) elif isinstance(batch[0], container_abcs.Sequence): return [[torch.from_numpy(sample) for sample in b] for b in batch] elif isinstance(batch[0], container_abcs.Mapping): for key in batch[0]: if key == 'coors' or key == 'feats': continue res[key] = collate_fn_([d[key] for d in batch]) return res res = collate_fn_(list_data) return res ================================================ FILE: dataset/simplify_dataset.py ================================================ import numpy as np import os import argparse parser = argparse.ArgumentParser() parser.add_argument('--dataset_root', default=None, required=True) def simplify_grasp_labels(root, save_path): """ original dataset grasp_label files have redundant data, We can significantly save the memory cost """ obj_names = list(range(88)) if not os.path.exists(save_path): os.makedirs(save_path) for i in obj_names: print('\nsimplifying object {}:'.format(i)) label = np.load(os.path.join(root, 'grasp_label', '{}_labels.npz'.format(str(i).zfill(3)))) # point_num = len(label['points']) print('original shape: ', label['points'].shape, label['offsets'].shape, label['scores'].shape) # if point_num > 4820: # idxs = np.random.choice(point_num, 4820, False) # points = label['points'][idxs] # offsets = label['offsets'][idxs] # scores = label['scores'][idxs] # print('Warning!!! down sample object {}'.format(i)) # else: points = label['points'] scores = label['scores'] offsets = label['offsets'] width = offsets[:, :, :, :, 2] print('after simplify, offset shape: ', points.shape, scores.shape, width.shape) np.savez(os.path.join(save_path, '{}_labels.npz'.format(str(i).zfill(3))), points=points, scores=scores, width=width) if __name__ == '__main__': cfgs = parser.parse_args() root = cfgs.dataset_root # set root and save path save_path = os.path.join(root, 'grasp_label_simplified') simplify_grasp_labels(root, save_path) ================================================ FILE: dataset/vis_graspness.py ================================================ import open3d as o3d import scipy.io as scio from PIL import Image import os import numpy as np import sys ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(ROOT_DIR) from utils.data_utils import get_workspace_mask, CameraInfo, create_point_cloud_from_depth_image data_path = '/media/bot/980A6F5E0A6F38801/datasets/graspnet/' scene_id = 'scene_0060' ann_id = '0000' camera_type = 'realsense' color = np.array(Image.open(os.path.join(data_path, 'scenes', scene_id, camera_type, 'rgb', ann_id + '.png')), dtype=np.float32) / 255.0 depth = np.array(Image.open(os.path.join(data_path, 'scenes', scene_id, camera_type, 'depth', ann_id + '.png'))) seg = np.array(Image.open(os.path.join(data_path, 'scenes', scene_id, camera_type, 'label', ann_id + '.png'))) meta = scio.loadmat(os.path.join(data_path, 'scenes', scene_id, camera_type, 'meta', ann_id + '.mat')) intrinsic = meta['intrinsic_matrix'] factor_depth = meta['factor_depth'] camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2], factor_depth) point_cloud = create_point_cloud_from_depth_image(depth, camera, organized=True) depth_mask = (depth > 0) camera_poses = np.load(os.path.join(data_path, 'scenes', scene_id, camera_type, 'camera_poses.npy')) align_mat = np.load(os.path.join(data_path, 'scenes', scene_id, camera_type, 'cam0_wrt_table.npy')) trans = np.dot(align_mat, camera_poses[int(ann_id)]) workspace_mask = get_workspace_mask(point_cloud, seg, trans=trans, organized=True, outlier=0.02) mask = (depth_mask & workspace_mask) point_cloud = point_cloud[mask] color = color[mask] seg = seg[mask] graspness_full = np.load(os.path.join(data_path, 'graspness', scene_id, camera_type, ann_id + '.npy')).squeeze() graspness_full[seg == 0] = 0. print('graspness full scene: ', graspness_full.shape, (graspness_full > 0.1).sum()) color[graspness_full > 0.1] = [0., 1., 0.] cloud = o3d.geometry.PointCloud() cloud.points = o3d.utility.Vector3dVector(point_cloud.astype(np.float32)) cloud.colors = o3d.utility.Vector3dVector(color.astype(np.float32)) o3d.visualization.draw_geometries([cloud]) ================================================ FILE: infer_vis_grasp.py ================================================ import os import sys import numpy as np import argparse from PIL import Image import time import scipy.io as scio import torch import open3d as o3d from graspnetAPI.graspnet_eval import GraspGroup ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.append(ROOT_DIR) sys.path.append(os.path.join(ROOT_DIR, 'utils')) from models.graspnet import GraspNet, pred_decode from dataset.graspnet_dataset import minkowski_collate_fn from collision_detector import ModelFreeCollisionDetector from data_utils import CameraInfo, create_point_cloud_from_depth_image, get_workspace_mask parser = argparse.ArgumentParser() parser.add_argument('--dataset_root', default='/data/datasets/graspnet') parser.add_argument('--checkpoint_path', default='/data/zibo/logs/graspness_kn.tar') parser.add_argument('--dump_dir', help='Dump dir to save outputs', default='/data/zibo/logs/') parser.add_argument('--seed_feat_dim', default=512, type=int, help='Point wise feature dim') parser.add_argument('--camera', default='kinect', help='Camera split [realsense/kinect]') parser.add_argument('--num_point', type=int, default=15000, help='Point Number [default: 15000]') parser.add_argument('--batch_size', type=int, default=1, help='Batch Size during inference [default: 1]') parser.add_argument('--voxel_size', type=float, default=0.005, help='Voxel Size for sparse convolution') parser.add_argument('--collision_thresh', type=float, default=-1, help='Collision Threshold in collision detection [default: 0.01]') parser.add_argument('--voxel_size_cd', type=float, default=0.01, help='Voxel Size for collision detection') parser.add_argument('--infer', action='store_true', default=False) parser.add_argument('--vis', action='store_true', default=False) parser.add_argument('--scene', type=str, default='0188') parser.add_argument('--index', type=str, default='0000') cfgs = parser.parse_args() # ------------------------------------------------------------------------- GLOBAL CONFIG BEG if not os.path.exists(cfgs.dump_dir): os.mkdir(cfgs.dump_dir) def data_process(): root = cfgs.dataset_root camera_type = cfgs.camera depth = np.array(Image.open(os.path.join(root, 'scenes', scene_id, camera_type, 'depth', index + '.png'))) seg = np.array(Image.open(os.path.join(root, 'scenes', scene_id, camera_type, 'label', index + '.png'))) meta = scio.loadmat(os.path.join(root, 'scenes', scene_id, camera_type, 'meta', index + '.mat')) try: intrinsic = meta['intrinsic_matrix'] factor_depth = meta['factor_depth'] except Exception as e: print(repr(e)) camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2], factor_depth) # generate cloud cloud = create_point_cloud_from_depth_image(depth, camera, organized=True) # get valid points depth_mask = (depth > 0) camera_poses = np.load(os.path.join(root, 'scenes', scene_id, camera_type, 'camera_poses.npy')) align_mat = np.load(os.path.join(root, 'scenes', scene_id, camera_type, 'cam0_wrt_table.npy')) trans = np.dot(align_mat, camera_poses[int(index)]) workspace_mask = get_workspace_mask(cloud, seg, trans=trans, organized=True, outlier=0.02) mask = (depth_mask & workspace_mask) cloud_masked = cloud[mask] # sample points random if len(cloud_masked) >= cfgs.num_point: idxs = np.random.choice(len(cloud_masked), cfgs.num_point, replace=False) else: idxs1 = np.arange(len(cloud_masked)) idxs2 = np.random.choice(len(cloud_masked), cfgs.num_point - len(cloud_masked), replace=True) idxs = np.concatenate([idxs1, idxs2], axis=0) cloud_sampled = cloud_masked[idxs] ret_dict = {'point_clouds': cloud_sampled.astype(np.float32), 'coors': cloud_sampled.astype(np.float32) / cfgs.voxel_size, 'feats': np.ones_like(cloud_sampled).astype(np.float32), } return ret_dict # Init datasets and dataloaders def my_worker_init_fn(worker_id): np.random.seed(np.random.get_state()[1][0] + worker_id) pass def inference(data_input): batch_data = minkowski_collate_fn([data_input]) net = GraspNet(seed_feat_dim=cfgs.seed_feat_dim, is_training=False) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device) # Load checkpoint checkpoint = torch.load(cfgs.checkpoint_path) net.load_state_dict(checkpoint['model_state_dict']) start_epoch = checkpoint['epoch'] print("-> loaded checkpoint %s (epoch: %d)" % (cfgs.checkpoint_path, start_epoch)) net.eval() tic = time.time() for key in batch_data: if 'list' in key: for i in range(len(batch_data[key])): for j in range(len(batch_data[key][i])): batch_data[key][i][j] = batch_data[key][i][j].to(device) else: batch_data[key] = batch_data[key].to(device) # Forward pass with torch.no_grad(): end_points = net(batch_data) grasp_preds = pred_decode(end_points) preds = grasp_preds[0].detach().cpu().numpy() gg = GraspGroup(preds) # collision detection if cfgs.collision_thresh > 0: cloud = data_input['point_clouds'] mfcdetector = ModelFreeCollisionDetector(cloud, voxel_size=cfgs.voxel_size_cd) collision_mask = mfcdetector.detect(gg, approach_dist=0.05, collision_thresh=cfgs.collision_thresh) gg = gg[~collision_mask] # save grasps save_dir = os.path.join(cfgs.dump_dir, scene_id, cfgs.camera) save_path = os.path.join(save_dir, cfgs.index + '.npy') if not os.path.exists(save_dir): os.makedirs(save_dir) gg.save_npy(save_path) toc = time.time() print('inference time: %fs' % (toc - tic)) if __name__ == '__main__': scene_id = 'scene_' + cfgs.scene index = cfgs.index data_dict = data_process() if cfgs.infer: inference(data_dict) if cfgs.vis: pc = data_dict['point_clouds'] gg = np.load(os.path.join(cfgs.dump_dir, scene_id, cfgs.camera, cfgs.index + '.npy')) gg = GraspGroup(gg) gg = gg.nms() gg = gg.sort_by_score() if gg.__len__() > 30: gg = gg[:30] grippers = gg.to_open3d_geometry_list() cloud = o3d.geometry.PointCloud() cloud.points = o3d.utility.Vector3dVector(pc.astype(np.float32)) o3d.visualization.draw_geometries([cloud, *grippers]) ================================================ FILE: knn/knn_modules.py ================================================ import unittest import gc import operator as op import functools import torch from torch.autograd import Variable, Function from knn_pytorch import knn_pytorch # import knn_pytorch def knn(ref, query, k=1): """ Compute k nearest neighbors for each query point. """ device = ref.device ref = ref.float().to(device) query = query.float().to(device) inds = torch.empty(query.shape[0], k, query.shape[2]).long().to(device) knn_pytorch.knn(ref, query, inds) return inds ================================================ FILE: knn/setup.py ================================================ #!/usr/bin/env python import glob import os import torch from setuptools import find_packages from setuptools import setup from torch.utils.cpp_extension import CUDA_HOME from torch.utils.cpp_extension import CppExtension from torch.utils.cpp_extension import CUDAExtension requirements = ["torch", "torchvision"] def get_extensions(): this_dir = os.path.dirname(os.path.abspath(__file__)) extensions_dir = os.path.join(this_dir, "src") main_file = glob.glob(os.path.join(extensions_dir, "*.cpp")) source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp")) source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu")) sources = main_file + source_cpu extension = CppExtension extra_compile_args = {"cxx": []} define_macros = [] if torch.cuda.is_available() and CUDA_HOME is not None: extension = CUDAExtension sources += source_cuda define_macros += [("WITH_CUDA", None)] extra_compile_args["nvcc"] = [ "-DCUDA_HAS_FP16=1", "-D__CUDA_NO_HALF_OPERATORS__", "-D__CUDA_NO_HALF_CONVERSIONS__", "-D__CUDA_NO_HALF2_OPERATORS__", ] sources = [os.path.join(extensions_dir, s) for s in sources] include_dirs = [extensions_dir] ext_modules = [ extension( "knn_pytorch.knn_pytorch", sources, include_dirs=include_dirs, define_macros=define_macros, extra_compile_args=extra_compile_args, ) ] return ext_modules setup( name="knn_pytorch", version="0.1", author="foolyc", url="https://github.com/foolyc/torchKNN", description="KNN implement in Pytorch 1.0 including both cpu version and gpu version", ext_modules=get_extensions(), cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}, ) ================================================ FILE: knn/src/cpu/knn_cpu.cpp ================================================ #include "cpu/vision.h" void knn_cpu(float* ref_dev, int ref_width, float* query_dev, int query_width, int height, int k, float* dist_dev, long* ind_dev, long* ind_buf) { // Compute all the distances for(int query_idx = 0;query_idx dist_dev[query_idx * ref_width + j + 1]) { temp_value = dist_dev[query_idx * ref_width + j]; dist_dev[query_idx * ref_width + j] = dist_dev[query_idx * ref_width + j + 1]; dist_dev[query_idx * ref_width + j + 1] = temp_value; temp_idx = ind_buf[j]; ind_buf[j] = ind_buf[j + 1]; ind_buf[j + 1] = temp_idx; } } for(int i = 0;i < k;i++) ind_dev[query_idx + i * query_width] = ind_buf[i]; #if DEBUG for(int i = 0;i < ref_width;i++) printf("%d, ", ind_buf[i]); printf("\n"); #endif } } ================================================ FILE: knn/src/cpu/vision.h ================================================ #pragma once #include void knn_cpu(float* ref_dev, int ref_width, float* query_dev, int query_width, int height, int k, float* dist_dev, long* ind_dev, long* ind_buf); ================================================ FILE: knn/src/cuda/knn.cu ================================================ /** Modifed version of knn-CUDA from https://github.com/vincentfpgarcia/kNN-CUDA * The modifications are * removed texture memory usage * removed split query KNN computation * added feature extraction with bilinear interpolation * * Last modified by Christopher B. Choy 12/23/2016 */ // Includes #include #include "cuda.h" #define IDX2D(i, j, dj) (dj * i + j) #define IDX3D(i, j, k, dj, dk) (IDX2D(IDX2D(i, j, dj), k, dk)) #define BLOCK 512 #define MAX_STREAMS 512 // Constants used by the program #define BLOCK_DIM 16 #define DEBUG 0 /** * Computes the distance between two matrix A (reference points) and * B (query points) containing respectively wA and wB points. * * @param A pointer on the matrix A * @param wA width of the matrix A = number of points in A * @param B pointer on the matrix B * @param wB width of the matrix B = number of points in B * @param dim dimension of points = height of matrices A and B * @param AB pointer on the matrix containing the wA*wB distances computed */ __global__ void cuComputeDistanceGlobal( float* A, int wA, float* B, int wB, int dim, float* AB){ // Declaration of the shared memory arrays As and Bs used to store the sub-matrix of A and B __shared__ float shared_A[BLOCK_DIM][BLOCK_DIM]; __shared__ float shared_B[BLOCK_DIM][BLOCK_DIM]; // Sub-matrix of A (begin, step, end) and Sub-matrix of B (begin, step) __shared__ int begin_A; __shared__ int begin_B; __shared__ int step_A; __shared__ int step_B; __shared__ int end_A; // Thread index int tx = threadIdx.x; int ty = threadIdx.y; // Other variables float tmp; float ssd = 0; // Loop parameters begin_A = BLOCK_DIM * blockIdx.y; begin_B = BLOCK_DIM * blockIdx.x; step_A = BLOCK_DIM * wA; step_B = BLOCK_DIM * wB; end_A = begin_A + (dim-1) * wA; // Conditions int cond0 = (begin_A + tx < wA); // used to write in shared memory int cond1 = (begin_B + tx < wB); // used to write in shared memory & to computations and to write in output matrix int cond2 = (begin_A + ty < wA); // used to computations and to write in output matrix // Loop over all the sub-matrices of A and B required to compute the block sub-matrix for (int a = begin_A, b = begin_B; a <= end_A; a += step_A, b += step_B) { // Load the matrices from device memory to shared memory; each thread loads one element of each matrix if (a/wA + ty < dim){ shared_A[ty][tx] = (cond0)? A[a + wA * ty + tx] : 0; shared_B[ty][tx] = (cond1)? B[b + wB * ty + tx] : 0; } else{ shared_A[ty][tx] = 0; shared_B[ty][tx] = 0; } // Synchronize to make sure the matrices are loaded __syncthreads(); // Compute the difference between the two matrixes; each thread computes one element of the block sub-matrix if (cond2 && cond1){ for (int k = 0; k < BLOCK_DIM; ++k){ tmp = shared_A[k][ty] - shared_B[k][tx]; ssd += tmp*tmp; } } // Synchronize to make sure that the preceding computation is done before loading two new sub-matrices of A and B in the next iteration __syncthreads(); } // Write the block sub-matrix to device memory; each thread writes one element if (cond2 && cond1) AB[(begin_A + ty) * wB + begin_B + tx] = ssd; } /** * Gathers k-th smallest distances for each column of the distance matrix in the top. * * @param dist distance matrix * @param ind index matrix * @param width width of the distance matrix and of the index matrix * @param height height of the distance matrix and of the index matrix * @param k number of neighbors to consider */ __global__ void cuInsertionSort(float *dist, long *ind, int width, int height, int k){ // Variables int l, i, j; float *p_dist; long *p_ind; float curr_dist, max_dist; long curr_row, max_row; unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x; if (xIndexcurr_dist){ i=a; break; } } for (j=l; j>i; j--){ p_dist[j*width] = p_dist[(j-1)*width]; p_ind[j*width] = p_ind[(j-1)*width]; } p_dist[i*width] = curr_dist; p_ind[i*width] = l+1; } else { p_ind[l*width] = l+1; } max_dist = p_dist[curr_row]; } // Part 2 : insert element in the k-th first lines max_row = (k-1)*width; for (l=k; lcurr_dist){ i=a; break; } } for (j=k-1; j>i; j--){ p_dist[j*width] = p_dist[(j-1)*width]; p_ind[j*width] = p_ind[(j-1)*width]; } p_dist[i*width] = curr_dist; p_ind[i*width] = l+1; max_dist = p_dist[max_row]; } } } } /** * Computes the square root of the first line (width-th first element) * of the distance matrix. * * @param dist distance matrix * @param width width of the distance matrix * @param k number of neighbors to consider */ __global__ void cuParallelSqrt(float *dist, int width, int k){ unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x; unsigned int yIndex = blockIdx.y * blockDim.y + threadIdx.y; if (xIndex>>(ref_dev, ref_nb, query_dev, query_nb, dim, dist_dev); // Kernel 2: Sort each column cuInsertionSort<<>>(dist_dev, ind_dev, query_nb, ref_nb, k); // Kernel 3: Compute square root of k first elements // cuParallelSqrt<<>>(dist_dev, query_nb, k); #if DEBUG unsigned int size_of_float = sizeof(float); unsigned long size_of_long = sizeof(long); float* dist_host = new float[query_nb * k]; long* idx_host = new long[query_nb * k]; // Memory copy of output from device to host cudaMemcpy(&dist_host[0], dist_dev, query_nb * k *size_of_float, cudaMemcpyDeviceToHost); cudaMemcpy(&idx_host[0], ind_dev, query_nb * k * size_of_long, cudaMemcpyDeviceToHost); int i = 0; for(i = 0; i < 100; i++){ printf("IDX[%d]: %d\n", i, (int)idx_host[i]); } #endif } ================================================ FILE: knn/src/cuda/vision.h ================================================ #pragma once #include #include void knn_device(float* ref_dev, int ref_width, float* query_dev, int query_width, int height, int k, float* dist_dev, long* ind_dev, cudaStream_t stream); ================================================ FILE: knn/src/knn.h ================================================ #pragma once #include "cpu/vision.h" #ifdef WITH_CUDA #include "cuda/vision.h" #include extern THCState *state; #endif int knn(at::Tensor& ref, at::Tensor& query, at::Tensor& idx) { // TODO check dimensions long batch, ref_nb, query_nb, dim, k; batch = ref.size(0); dim = ref.size(1); k = idx.size(1); ref_nb = ref.size(2); query_nb = query.size(2); float *ref_dev = ref.data(); float *query_dev = query.data(); long *idx_dev = idx.data(); if (ref.type().is_cuda()) { #ifdef WITH_CUDA // TODO raise error if not compiled with CUDA float *dist_dev = (float*)THCudaMalloc(state, ref_nb * query_nb * sizeof(float)); for (int b = 0; b < batch; b++) { // knn_device(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k, // dist_dev, idx_dev + b * k * query_nb, THCState_getCurrentStream(state)); knn_device(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k, dist_dev, idx_dev + b * k * query_nb, c10::cuda::getCurrentCUDAStream()); } THCudaFree(state, dist_dev); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { printf("error in knn: %s\n", cudaGetErrorString(err)); THError("aborting"); } return 1; #else AT_ERROR("Not compiled with GPU support"); #endif } float *dist_dev = (float*)malloc(ref_nb * query_nb * sizeof(float)); long *ind_buf = (long*)malloc(ref_nb * sizeof(long)); for (int b = 0; b < batch; b++) { knn_cpu(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k, dist_dev, idx_dev + b * k * query_nb, ind_buf); } free(dist_dev); free(ind_buf); return 1; } ================================================ FILE: knn/src/vision.cpp ================================================ #include "knn.h" PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("knn", &knn, "k-nearest neighbors"); } ================================================ FILE: models/backbone_resunet14.py ================================================ import MinkowskiEngine as ME from MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck from models.resnet import ResNetBase class MinkUNetBase(ResNetBase): BLOCK = None PLANES = None DILATIONS = (1, 1, 1, 1, 1, 1, 1, 1) LAYERS = (2, 2, 2, 2, 2, 2, 2, 2) PLANES = (32, 64, 128, 256, 256, 128, 96, 96) INIT_DIM = 32 OUT_TENSOR_STRIDE = 1 # To use the model, must call initialize_coords before forward pass. # Once data is processed, call clear to reset the model before calling # initialize_coords def __init__(self, in_channels, out_channels, D=3): ResNetBase.__init__(self, in_channels, out_channels, D) def network_initialization(self, in_channels, out_channels, D): # Output of the first conv concated to conv6 self.inplanes = self.INIT_DIM self.conv0p1s1 = ME.MinkowskiConvolution( in_channels, self.inplanes, kernel_size=5, dimension=D) self.bn0 = ME.MinkowskiBatchNorm(self.inplanes) self.conv1p1s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn1 = ME.MinkowskiBatchNorm(self.inplanes) self.block1 = self._make_layer(self.BLOCK, self.PLANES[0], self.LAYERS[0]) self.conv2p2s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn2 = ME.MinkowskiBatchNorm(self.inplanes) self.block2 = self._make_layer(self.BLOCK, self.PLANES[1], self.LAYERS[1]) self.conv3p4s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn3 = ME.MinkowskiBatchNorm(self.inplanes) self.block3 = self._make_layer(self.BLOCK, self.PLANES[2], self.LAYERS[2]) self.conv4p8s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn4 = ME.MinkowskiBatchNorm(self.inplanes) self.block4 = self._make_layer(self.BLOCK, self.PLANES[3], self.LAYERS[3]) self.convtr4p16s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[4], kernel_size=2, stride=2, dimension=D) self.bntr4 = ME.MinkowskiBatchNorm(self.PLANES[4]) self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion self.block5 = self._make_layer(self.BLOCK, self.PLANES[4], self.LAYERS[4]) self.convtr5p8s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[5], kernel_size=2, stride=2, dimension=D) self.bntr5 = ME.MinkowskiBatchNorm(self.PLANES[5]) self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion self.block6 = self._make_layer(self.BLOCK, self.PLANES[5], self.LAYERS[5]) self.convtr6p4s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[6], kernel_size=2, stride=2, dimension=D) self.bntr6 = ME.MinkowskiBatchNorm(self.PLANES[6]) self.inplanes = self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion self.block7 = self._make_layer(self.BLOCK, self.PLANES[6], self.LAYERS[6]) self.convtr7p2s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[7], kernel_size=2, stride=2, dimension=D) self.bntr7 = ME.MinkowskiBatchNorm(self.PLANES[7]) self.inplanes = self.PLANES[7] + self.INIT_DIM self.block8 = self._make_layer(self.BLOCK, self.PLANES[7], self.LAYERS[7]) self.final = ME.MinkowskiConvolution( self.PLANES[7] * self.BLOCK.expansion, out_channels, kernel_size=1, bias=True, dimension=D) self.relu = ME.MinkowskiReLU(inplace=True) def forward(self, x): out = self.conv0p1s1(x) out = self.bn0(out) out_p1 = self.relu(out) out = self.conv1p1s2(out_p1) out = self.bn1(out) out = self.relu(out) out_b1p2 = self.block1(out) out = self.conv2p2s2(out_b1p2) out = self.bn2(out) out = self.relu(out) out_b2p4 = self.block2(out) out = self.conv3p4s2(out_b2p4) out = self.bn3(out) out = self.relu(out) out_b3p8 = self.block3(out) # tensor_stride=16 out = self.conv4p8s2(out_b3p8) out = self.bn4(out) out = self.relu(out) out = self.block4(out) # tensor_stride=8 out = self.convtr4p16s2(out) out = self.bntr4(out) out = self.relu(out) out = ME.cat(out, out_b3p8) out = self.block5(out) # tensor_stride=4 out = self.convtr5p8s2(out) out = self.bntr5(out) out = self.relu(out) out = ME.cat(out, out_b2p4) out = self.block6(out) # tensor_stride=2 out = self.convtr6p4s2(out) out = self.bntr6(out) out = self.relu(out) out = ME.cat(out, out_b1p2) out = self.block7(out) # tensor_stride=1 out = self.convtr7p2s2(out) out = self.bntr7(out) out = self.relu(out) out = ME.cat(out, out_p1) out = self.block8(out) return self.final(out) class MinkUNet14(MinkUNetBase): BLOCK = BasicBlock LAYERS = (1, 1, 1, 1, 1, 1, 1, 1) class MinkUNet18(MinkUNetBase): BLOCK = BasicBlock LAYERS = (2, 2, 2, 2, 2, 2, 2, 2) class MinkUNet34(MinkUNetBase): BLOCK = BasicBlock LAYERS = (2, 3, 4, 6, 2, 2, 2, 2) class MinkUNet50(MinkUNetBase): BLOCK = Bottleneck LAYERS = (2, 3, 4, 6, 2, 2, 2, 2) class MinkUNet101(MinkUNetBase): BLOCK = Bottleneck LAYERS = (2, 3, 4, 23, 2, 2, 2, 2) class MinkUNet14A(MinkUNet14): PLANES = (32, 64, 128, 256, 128, 128, 96, 96) class MinkUNet14B(MinkUNet14): PLANES = (32, 64, 128, 256, 128, 128, 128, 128) class MinkUNet14C(MinkUNet14): PLANES = (32, 64, 128, 256, 192, 192, 128, 128) class MinkUNet14Dori(MinkUNet14): PLANES = (32, 64, 128, 256, 384, 384, 384, 384) class MinkUNet14E(MinkUNet14): PLANES = (32, 64, 128, 256, 384, 384, 384, 384) class MinkUNet14D(MinkUNet14): PLANES = (32, 64, 128, 256, 192, 192, 192, 192) class MinkUNet18A(MinkUNet18): PLANES = (32, 64, 128, 256, 128, 128, 96, 96) class MinkUNet18B(MinkUNet18): PLANES = (32, 64, 128, 256, 128, 128, 128, 128) class MinkUNet18D(MinkUNet18): PLANES = (32, 64, 128, 256, 384, 384, 384, 384) class MinkUNet34A(MinkUNet34): PLANES = (32, 64, 128, 256, 256, 128, 64, 64) class MinkUNet34B(MinkUNet34): PLANES = (32, 64, 128, 256, 256, 128, 64, 32) class MinkUNet34C(MinkUNet34): PLANES = (32, 64, 128, 256, 256, 128, 96, 96) ================================================ FILE: models/graspnet.py ================================================ """ GraspNet baseline model definition. Author: chenxi-wang """ import os import sys import numpy as np import torch import torch.nn as nn import MinkowskiEngine as ME BASE_DIR = os.path.dirname(os.path.abspath(__file__)) ROOT_DIR = os.path.dirname(BASE_DIR) sys.path.append(ROOT_DIR) from models.backbone_resunet14 import MinkUNet14D from models.modules import ApproachNet, GraspableNet, CloudCrop, SWADNet from loss_utils import GRASP_MAX_WIDTH, NUM_VIEW, NUM_ANGLE, NUM_DEPTH, GRASPNESS_THRESHOLD, M_POINT from label_generation import process_grasp_labels, match_grasp_view_and_label, batch_viewpoint_params_to_matrix from pointnet2.pointnet2_utils import furthest_point_sample, gather_operation class GraspNet(nn.Module): def __init__(self, cylinder_radius=0.05, seed_feat_dim=512, is_training=True): super().__init__() self.is_training = is_training self.seed_feature_dim = seed_feat_dim self.num_depth = NUM_DEPTH self.num_angle = NUM_ANGLE self.M_points = M_POINT self.num_view = NUM_VIEW self.backbone = MinkUNet14D(in_channels=3, out_channels=self.seed_feature_dim, D=3) self.graspable = GraspableNet(seed_feature_dim=self.seed_feature_dim) self.rotation = ApproachNet(self.num_view, seed_feature_dim=self.seed_feature_dim, is_training=self.is_training) self.crop = CloudCrop(nsample=16, cylinder_radius=cylinder_radius, seed_feature_dim=self.seed_feature_dim) self.swad = SWADNet(num_angle=self.num_angle, num_depth=self.num_depth) def forward(self, end_points): seed_xyz = end_points['point_clouds'] # use all sampled point cloud, B*Ns*3 B, point_num, _ = seed_xyz.shape # batch _size # point-wise features coordinates_batch = end_points['coors'] features_batch = end_points['feats'] mink_input = ME.SparseTensor(features_batch, coordinates=coordinates_batch) seed_features = self.backbone(mink_input).F seed_features = seed_features[end_points['quantize2original']].view(B, point_num, -1).transpose(1, 2) end_points = self.graspable(seed_features, end_points) seed_features_flipped = seed_features.transpose(1, 2) # B*Ns*feat_dim objectness_score = end_points['objectness_score'] graspness_score = end_points['graspness_score'].squeeze(1) objectness_pred = torch.argmax(objectness_score, 1) objectness_mask = (objectness_pred == 1) graspness_mask = graspness_score > GRASPNESS_THRESHOLD graspable_mask = objectness_mask & graspness_mask seed_features_graspable = [] seed_xyz_graspable = [] graspable_num_batch = 0. for i in range(B): cur_mask = graspable_mask[i] graspable_num_batch += cur_mask.sum() cur_feat = seed_features_flipped[i][cur_mask] # Ns*feat_dim cur_seed_xyz = seed_xyz[i][cur_mask] # Ns*3 cur_seed_xyz = cur_seed_xyz.unsqueeze(0) # 1*Ns*3 fps_idxs = furthest_point_sample(cur_seed_xyz, self.M_points) cur_seed_xyz_flipped = cur_seed_xyz.transpose(1, 2).contiguous() # 1*3*Ns cur_seed_xyz = gather_operation(cur_seed_xyz_flipped, fps_idxs).transpose(1, 2).squeeze(0).contiguous() # Ns*3 cur_feat_flipped = cur_feat.unsqueeze(0).transpose(1, 2).contiguous() # 1*feat_dim*Ns cur_feat = gather_operation(cur_feat_flipped, fps_idxs).squeeze(0).contiguous() # feat_dim*Ns seed_features_graspable.append(cur_feat) seed_xyz_graspable.append(cur_seed_xyz) seed_xyz_graspable = torch.stack(seed_xyz_graspable, 0) # B*Ns*3 seed_features_graspable = torch.stack(seed_features_graspable) # B*feat_dim*Ns end_points['xyz_graspable'] = seed_xyz_graspable end_points['graspable_count_stage1'] = graspable_num_batch / B end_points, res_feat = self.rotation(seed_features_graspable, end_points) seed_features_graspable = seed_features_graspable + res_feat if self.is_training: end_points = process_grasp_labels(end_points) grasp_top_views_rot, end_points = match_grasp_view_and_label(end_points) else: grasp_top_views_rot = end_points['grasp_top_view_rot'] group_features = self.crop(seed_xyz_graspable.contiguous(), seed_features_graspable.contiguous(), grasp_top_views_rot) end_points = self.swad(group_features, end_points) return end_points def pred_decode(end_points): batch_size = len(end_points['point_clouds']) grasp_preds = [] for i in range(batch_size): grasp_center = end_points['xyz_graspable'][i].float() grasp_score = end_points['grasp_score_pred'][i].float() grasp_score = grasp_score.view(M_POINT, NUM_ANGLE*NUM_DEPTH) grasp_score, grasp_score_inds = torch.max(grasp_score, -1) # [M_POINT] grasp_score = grasp_score.view(-1, 1) grasp_angle = (grasp_score_inds // NUM_DEPTH) * np.pi / 12 grasp_depth = (grasp_score_inds % NUM_DEPTH + 1) * 0.01 grasp_depth = grasp_depth.view(-1, 1) grasp_width = 1.2 * end_points['grasp_width_pred'][i] / 10. grasp_width = grasp_width.view(M_POINT, NUM_ANGLE*NUM_DEPTH) grasp_width = torch.gather(grasp_width, 1, grasp_score_inds.view(-1, 1)) grasp_width = torch.clamp(grasp_width, min=0., max=GRASP_MAX_WIDTH) approaching = -end_points['grasp_top_view_xyz'][i].float() grasp_rot = batch_viewpoint_params_to_matrix(approaching, grasp_angle) grasp_rot = grasp_rot.view(M_POINT, 9) # merge preds grasp_height = 0.02 * torch.ones_like(grasp_score) obj_ids = -1 * torch.ones_like(grasp_score) grasp_preds.append( torch.cat([grasp_score, grasp_width, grasp_height, grasp_depth, grasp_rot, grasp_center, obj_ids], axis=-1)) return grasp_preds ================================================ FILE: models/loss.py ================================================ import torch.nn as nn import torch def get_loss(end_points): objectness_loss, end_points = compute_objectness_loss(end_points) graspness_loss, end_points = compute_graspness_loss(end_points) view_loss, end_points = compute_view_graspness_loss(end_points) score_loss, end_points = compute_score_loss(end_points) width_loss, end_points = compute_width_loss(end_points) loss = objectness_loss + 10 * graspness_loss + 100 * view_loss + 15 * score_loss + 10 * width_loss end_points['loss/overall_loss'] = loss return loss, end_points def compute_objectness_loss(end_points): criterion = nn.CrossEntropyLoss(reduction='mean') objectness_score = end_points['objectness_score'] objectness_label = end_points['objectness_label'] loss = criterion(objectness_score, objectness_label) end_points['loss/stage1_objectness_loss'] = loss objectness_pred = torch.argmax(objectness_score, 1) end_points['stage1_objectness_acc'] = (objectness_pred == objectness_label.long()).float().mean() end_points['stage1_objectness_prec'] = (objectness_pred == objectness_label.long())[ objectness_pred == 1].float().mean() end_points['stage1_objectness_recall'] = (objectness_pred == objectness_label.long())[ objectness_label == 1].float().mean() return loss, end_points def compute_graspness_loss(end_points): criterion = nn.SmoothL1Loss(reduction='none') graspness_score = end_points['graspness_score'].squeeze(1) graspness_label = end_points['graspness_label'].squeeze(-1) loss_mask = end_points['objectness_label'].bool() loss = criterion(graspness_score, graspness_label) loss = loss[loss_mask] loss = loss.mean() graspness_score_c = graspness_score.detach().clone()[loss_mask] graspness_label_c = graspness_label.detach().clone()[loss_mask] graspness_score_c = torch.clamp(graspness_score_c, 0., 0.99) graspness_label_c = torch.clamp(graspness_label_c, 0., 0.99) rank_error = (torch.abs(torch.trunc(graspness_score_c * 20) - torch.trunc(graspness_label_c * 20)) / 20.).mean() end_points['stage1_graspness_acc_rank_error'] = rank_error end_points['loss/stage1_graspness_loss'] = loss return loss, end_points def compute_view_graspness_loss(end_points): criterion = nn.SmoothL1Loss(reduction='mean') view_score = end_points['view_score'] view_label = end_points['batch_grasp_view_graspness'] loss = criterion(view_score, view_label) end_points['loss/stage2_view_loss'] = loss return loss, end_points def compute_score_loss(end_points): criterion = nn.SmoothL1Loss(reduction='mean') grasp_score_pred = end_points['grasp_score_pred'] grasp_score_label = end_points['batch_grasp_score'] loss = criterion(grasp_score_pred, grasp_score_label) end_points['loss/stage3_score_loss'] = loss return loss, end_points def compute_width_loss(end_points): criterion = nn.SmoothL1Loss(reduction='none') grasp_width_pred = end_points['grasp_width_pred'] grasp_width_label = end_points['batch_grasp_width'] * 10 loss = criterion(grasp_width_pred, grasp_width_label) grasp_score_label = end_points['batch_grasp_score'] loss_mask = grasp_score_label > 0 loss = loss[loss_mask].mean() end_points['loss/stage3_width_loss'] = loss return loss, end_points ================================================ FILE: models/modules.py ================================================ import os import sys import torch import torch.nn as nn import torch.nn.functional as F BASE_DIR = os.path.dirname(os.path.abspath(__file__)) ROOT_DIR = os.path.dirname(BASE_DIR) sys.path.append(ROOT_DIR) import pointnet2.pytorch_utils as pt_utils from pointnet2.pointnet2_utils import CylinderQueryAndGroup from loss_utils import generate_grasp_views, batch_viewpoint_params_to_matrix class GraspableNet(nn.Module): def __init__(self, seed_feature_dim): super().__init__() self.in_dim = seed_feature_dim self.conv_graspable = nn.Conv1d(self.in_dim, 3, 1) def forward(self, seed_features, end_points): graspable_score = self.conv_graspable(seed_features) # (B, 3, num_seed) end_points['objectness_score'] = graspable_score[:, :2] end_points['graspness_score'] = graspable_score[:, 2] return end_points class ApproachNet(nn.Module): def __init__(self, num_view, seed_feature_dim, is_training=True): super().__init__() self.num_view = num_view self.in_dim = seed_feature_dim self.is_training = is_training self.conv1 = nn.Conv1d(self.in_dim, self.in_dim, 1) self.conv2 = nn.Conv1d(self.in_dim, self.num_view, 1) def forward(self, seed_features, end_points): B, _, num_seed = seed_features.size() res_features = F.relu(self.conv1(seed_features), inplace=True) features = self.conv2(res_features) view_score = features.transpose(1, 2).contiguous() # (B, num_seed, num_view) end_points['view_score'] = view_score if self.is_training: # normalize view graspness score to 0~1 view_score_ = view_score.clone().detach() view_score_max, _ = torch.max(view_score_, dim=2) view_score_min, _ = torch.min(view_score_, dim=2) view_score_max = view_score_max.unsqueeze(-1).expand(-1, -1, self.num_view) view_score_min = view_score_min.unsqueeze(-1).expand(-1, -1, self.num_view) view_score_ = (view_score_ - view_score_min) / (view_score_max - view_score_min + 1e-8) top_view_inds = [] for i in range(B): top_view_inds_batch = torch.multinomial(view_score_[i], 1, replacement=False) top_view_inds.append(top_view_inds_batch) top_view_inds = torch.stack(top_view_inds, dim=0).squeeze(-1) # B, num_seed else: _, top_view_inds = torch.max(view_score, dim=2) # (B, num_seed) top_view_inds_ = top_view_inds.view(B, num_seed, 1, 1).expand(-1, -1, -1, 3).contiguous() template_views = generate_grasp_views(self.num_view).to(features.device) # (num_view, 3) template_views = template_views.view(1, 1, self.num_view, 3).expand(B, num_seed, -1, -1).contiguous() vp_xyz = torch.gather(template_views, 2, top_view_inds_).squeeze(2) # (B, num_seed, 3) vp_xyz_ = vp_xyz.view(-1, 3) batch_angle = torch.zeros(vp_xyz_.size(0), dtype=vp_xyz.dtype, device=vp_xyz.device) vp_rot = batch_viewpoint_params_to_matrix(-vp_xyz_, batch_angle).view(B, num_seed, 3, 3) end_points['grasp_top_view_xyz'] = vp_xyz end_points['grasp_top_view_rot'] = vp_rot end_points['grasp_top_view_inds'] = top_view_inds return end_points, res_features class CloudCrop(nn.Module): def __init__(self, nsample, seed_feature_dim, cylinder_radius=0.05, hmin=-0.02, hmax=0.04): super().__init__() self.nsample = nsample self.in_dim = seed_feature_dim self.cylinder_radius = cylinder_radius mlps = [3 + self.in_dim, 256, 256] # use xyz, so plus 3 self.grouper = CylinderQueryAndGroup(radius=cylinder_radius, hmin=hmin, hmax=hmax, nsample=nsample, use_xyz=True, normalize_xyz=True) self.mlps = pt_utils.SharedMLP(mlps, bn=True) def forward(self, seed_xyz_graspable, seed_features_graspable, vp_rot): grouped_feature = self.grouper(seed_xyz_graspable, seed_xyz_graspable, vp_rot, seed_features_graspable) # B*3 + feat_dim*M*K new_features = self.mlps(grouped_feature) # (batch_size, mlps[-1], M, K) new_features = F.max_pool2d(new_features, kernel_size=[1, new_features.size(3)]) # (batch_size, mlps[-1], M, 1) new_features = new_features.squeeze(-1) # (batch_size, mlps[-1], M) return new_features class SWADNet(nn.Module): def __init__(self, num_angle, num_depth): super().__init__() self.num_angle = num_angle self.num_depth = num_depth self.conv1 = nn.Conv1d(256, 256, 1) # input feat dim need to be consistent with CloudCrop module self.conv_swad = nn.Conv1d(256, 2*num_angle*num_depth, 1) def forward(self, vp_features, end_points): B, _, num_seed = vp_features.size() vp_features = F.relu(self.conv1(vp_features), inplace=True) vp_features = self.conv_swad(vp_features) vp_features = vp_features.view(B, 2, self.num_angle, self.num_depth, num_seed) vp_features = vp_features.permute(0, 1, 4, 2, 3) # split prediction end_points['grasp_score_pred'] = vp_features[:, 0] # B * num_seed * num angle * num_depth end_points['grasp_width_pred'] = vp_features[:, 1] return end_points ================================================ FILE: models/resnet.py ================================================ import torch.nn as nn try: import open3d as o3d except ImportError: raise ImportError("Please install open3d with `pip install open3d`.") import MinkowskiEngine as ME from MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck class ResNetBase(nn.Module): BLOCK = None LAYERS = () INIT_DIM = 64 PLANES = (64, 128, 256, 512) def __init__(self, in_channels, out_channels, D=3): nn.Module.__init__(self) self.D = D assert self.BLOCK is not None self.network_initialization(in_channels, out_channels, D) self.weight_initialization() def network_initialization(self, in_channels, out_channels, D): self.inplanes = self.INIT_DIM self.conv1 = nn.Sequential( ME.MinkowskiConvolution( in_channels, self.inplanes, kernel_size=3, stride=2, dimension=D ), ME.MinkowskiInstanceNorm(self.inplanes), ME.MinkowskiReLU(inplace=True), ME.MinkowskiMaxPooling(kernel_size=2, stride=2, dimension=D), ) self.layer1 = self._make_layer( self.BLOCK, self.PLANES[0], self.LAYERS[0], stride=2 ) self.layer2 = self._make_layer( self.BLOCK, self.PLANES[1], self.LAYERS[1], stride=2 ) self.layer3 = self._make_layer( self.BLOCK, self.PLANES[2], self.LAYERS[2], stride=2 ) self.layer4 = self._make_layer( self.BLOCK, self.PLANES[3], self.LAYERS[3], stride=2 ) self.conv5 = nn.Sequential( ME.MinkowskiDropout(), ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=3, stride=3, dimension=D ), ME.MinkowskiInstanceNorm(self.inplanes), ME.MinkowskiGELU(), ) self.glob_pool = ME.MinkowskiGlobalMaxPooling() self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True) def weight_initialization(self): for m in self.modules(): if isinstance(m, ME.MinkowskiConvolution): ME.utils.kaiming_normal_(m.kernel, mode="fan_out", nonlinearity="relu") if isinstance(m, ME.MinkowskiBatchNorm): nn.init.constant_(m.bn.weight, 1) nn.init.constant_(m.bn.bias, 0) def _make_layer(self, block, planes, blocks, stride=1, dilation=1, bn_momentum=0.1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( ME.MinkowskiConvolution( self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, dimension=self.D, ), ME.MinkowskiBatchNorm(planes * block.expansion), ) layers = [] layers.append( block( self.inplanes, planes, stride=stride, dilation=dilation, downsample=downsample, dimension=self.D, ) ) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append( block( self.inplanes, planes, stride=1, dilation=dilation, dimension=self.D ) ) return nn.Sequential(*layers) def forward(self, x: ME.SparseTensor): x = self.conv1(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.conv5(x) x = self.glob_pool(x) return self.final(x) class ResNet14(ResNetBase): BLOCK = BasicBlock LAYERS = (1, 1, 1, 1) class ResNet18(ResNetBase): BLOCK = BasicBlock LAYERS = (2, 2, 2, 2) class ResNet34(ResNetBase): BLOCK = BasicBlock LAYERS = (3, 4, 6, 3) class ResNet50(ResNetBase): BLOCK = Bottleneck LAYERS = (3, 4, 6, 3) class ResNet101(ResNetBase): BLOCK = Bottleneck LAYERS = (3, 4, 23, 3) class ResFieldNetBase(ResNetBase): def network_initialization(self, in_channels, out_channels, D): field_ch = 32 field_ch2 = 64 self.field_network = nn.Sequential( ME.MinkowskiSinusoidal(in_channels, field_ch), ME.MinkowskiBatchNorm(field_ch), ME.MinkowskiReLU(inplace=True), ME.MinkowskiLinear(field_ch, field_ch), ME.MinkowskiBatchNorm(field_ch), ME.MinkowskiReLU(inplace=True), ME.MinkowskiToSparseTensor(), ) self.field_network2 = nn.Sequential( ME.MinkowskiSinusoidal(field_ch + in_channels, field_ch2), ME.MinkowskiBatchNorm(field_ch2), ME.MinkowskiReLU(inplace=True), ME.MinkowskiLinear(field_ch2, field_ch2), ME.MinkowskiBatchNorm(field_ch2), ME.MinkowskiReLU(inplace=True), ME.MinkowskiToSparseTensor(), ) ResNetBase.network_initialization(self, field_ch2, out_channels, D) def forward(self, x: ME.TensorField): otensor = self.field_network(x) otensor2 = self.field_network2(otensor.cat_slice(x)) return ResNetBase.forward(self, otensor2) class ResFieldNet14(ResFieldNetBase): BLOCK = BasicBlock LAYERS = (1, 1, 1, 1) class ResFieldNet18(ResFieldNetBase): BLOCK = BasicBlock LAYERS = (2, 2, 2, 2) class ResFieldNet34(ResFieldNetBase): BLOCK = BasicBlock LAYERS = (3, 4, 6, 3) class ResFieldNet50(ResFieldNetBase): BLOCK = Bottleneck LAYERS = (3, 4, 6, 3) class ResFieldNet101(ResFieldNetBase): BLOCK = Bottleneck LAYERS = (3, 4, 23, 3) ================================================ FILE: pointnet2/_ext_src/include/ball_query.h ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #pragma once #include at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, const int nsample); ================================================ FILE: pointnet2/_ext_src/include/cuda_utils.h ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #ifndef _CUDA_UTILS_H #define _CUDA_UTILS_H #include #include #include #include #include #include #define TOTAL_THREADS 512 inline int opt_n_threads(int work_size) { const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); return max(min(1 << pow_2, TOTAL_THREADS), 1); } inline dim3 opt_block_config(int x, int y) { const int x_threads = opt_n_threads(x); const int y_threads = max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); dim3 block_config(x_threads, y_threads, 1); return block_config; } #define CUDA_CHECK_ERRORS() \ do { \ cudaError_t err = cudaGetLastError(); \ if (cudaSuccess != err) { \ fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ __FILE__); \ exit(-1); \ } \ } while (0) #endif ================================================ FILE: pointnet2/_ext_src/include/cylinder_query.h ================================================ // Author: chenxi-wang #pragma once #include at::Tensor cylinder_query(at::Tensor new_xyz, at::Tensor xyz, at::Tensor rot, const float radius, const float hmin, const float hmax, const int nsample); ================================================ FILE: pointnet2/_ext_src/include/group_points.h ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #pragma once #include at::Tensor group_points(at::Tensor points, at::Tensor idx); at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); ================================================ FILE: pointnet2/_ext_src/include/interpolate.h ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #pragma once #include #include std::vector three_nn(at::Tensor unknowns, at::Tensor knows); at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, at::Tensor weight); at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, at::Tensor weight, const int m); ================================================ FILE: pointnet2/_ext_src/include/sampling.h ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #pragma once #include at::Tensor gather_points(at::Tensor points, at::Tensor idx); at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); ================================================ FILE: pointnet2/_ext_src/include/utils.h ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #pragma once #include #include #define CHECK_CUDA(x) \ do { \ TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ } while (0) #define CHECK_CONTIGUOUS(x) \ do { \ TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ } while (0) #define CHECK_IS_INT(x) \ do { \ TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \ #x " must be an int tensor"); \ } while (0) #define CHECK_IS_FLOAT(x) \ do { \ TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \ #x " must be a float tensor"); \ } while (0) ================================================ FILE: pointnet2/_ext_src/src/ball_query.cpp ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include "ball_query.h" #include "utils.h" void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, int nsample, const float *new_xyz, const float *xyz, int *idx); at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, const int nsample) { CHECK_CONTIGUOUS(new_xyz); CHECK_CONTIGUOUS(xyz); CHECK_IS_FLOAT(new_xyz); CHECK_IS_FLOAT(xyz); if (new_xyz.type().is_cuda()) { CHECK_CUDA(xyz); } at::Tensor idx = torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, at::device(new_xyz.device()).dtype(at::ScalarType::Int)); if (new_xyz.type().is_cuda()) { query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), radius, nsample, new_xyz.data(), xyz.data(), idx.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return idx; } ================================================ FILE: pointnet2/_ext_src/src/ball_query_gpu.cu ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include #include #include #include "cuda_utils.h" // input: new_xyz(b, m, 3) xyz(b, n, 3) // output: idx(b, m, nsample) __global__ void query_ball_point_kernel(int b, int n, int m, float radius, int nsample, const float *__restrict__ new_xyz, const float *__restrict__ xyz, int *__restrict__ idx) { int batch_index = blockIdx.x; xyz += batch_index * n * 3; new_xyz += batch_index * m * 3; idx += m * nsample * batch_index; int index = threadIdx.x; int stride = blockDim.x; float radius2 = radius * radius; for (int j = index; j < m; j += stride) { float new_x = new_xyz[j * 3 + 0]; float new_y = new_xyz[j * 3 + 1]; float new_z = new_xyz[j * 3 + 2]; for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { float x = xyz[k * 3 + 0]; float y = xyz[k * 3 + 1]; float z = xyz[k * 3 + 2]; float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); if (d2 < radius2) { if (cnt == 0) { for (int l = 0; l < nsample; ++l) { idx[j * nsample + l] = k; } } idx[j * nsample + cnt] = k; ++cnt; } } } } void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, int nsample, const float *new_xyz, const float *xyz, int *idx) { cudaStream_t stream = at::cuda::getCurrentCUDAStream(); query_ball_point_kernel<<>>( b, n, m, radius, nsample, new_xyz, xyz, idx); CUDA_CHECK_ERRORS(); } ================================================ FILE: pointnet2/_ext_src/src/bindings.cpp ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include "ball_query.h" #include "group_points.h" #include "interpolate.h" #include "sampling.h" #include "cylinder_query.h" PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("gather_points", &gather_points); m.def("gather_points_grad", &gather_points_grad); m.def("furthest_point_sampling", &furthest_point_sampling); m.def("three_nn", &three_nn); m.def("three_interpolate", &three_interpolate); m.def("three_interpolate_grad", &three_interpolate_grad); m.def("ball_query", &ball_query); m.def("group_points", &group_points); m.def("group_points_grad", &group_points_grad); m.def("cylinder_query", &cylinder_query); } ================================================ FILE: pointnet2/_ext_src/src/cylinder_query.cpp ================================================ // Author: chenxi-wang #include "cylinder_query.h" #include "utils.h" void query_cylinder_point_kernel_wrapper(int b, int n, int m, float radius, float hmin, float hmax, int nsample, const float *new_xyz, const float *xyz, const float *rot, int *idx); at::Tensor cylinder_query(at::Tensor new_xyz, at::Tensor xyz, at::Tensor rot, const float radius, const float hmin, const float hmax, const int nsample) { CHECK_CONTIGUOUS(new_xyz); CHECK_CONTIGUOUS(xyz); CHECK_CONTIGUOUS(rot); CHECK_IS_FLOAT(new_xyz); CHECK_IS_FLOAT(xyz); CHECK_IS_FLOAT(rot); if (new_xyz.type().is_cuda()) { CHECK_CUDA(xyz); CHECK_CUDA(rot); } at::Tensor idx = torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, at::device(new_xyz.device()).dtype(at::ScalarType::Int)); if (new_xyz.type().is_cuda()) { query_cylinder_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), radius, hmin, hmax, nsample, new_xyz.data(), xyz.data(), rot.data(), idx.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return idx; } ================================================ FILE: pointnet2/_ext_src/src/cylinder_query_gpu.cu ================================================ // Author: chenxi-wang #include #include #include #include "cuda_utils.h" __global__ void query_cylinder_point_kernel(int b, int n, int m, float radius, float hmin, float hmax, int nsample, const float *__restrict__ new_xyz, const float *__restrict__ xyz, const float *__restrict__ rot, int *__restrict__ idx) { int batch_index = blockIdx.x; xyz += batch_index * n * 3; new_xyz += batch_index * m * 3; rot += batch_index * m * 9; idx += m * nsample * batch_index; int index = threadIdx.x; int stride = blockDim.x; float radius2 = radius * radius; for (int j = index; j < m; j += stride) { float new_x = new_xyz[j * 3 + 0]; float new_y = new_xyz[j * 3 + 1]; float new_z = new_xyz[j * 3 + 2]; float r0 = rot[j * 9 + 0]; float r1 = rot[j * 9 + 1]; float r2 = rot[j * 9 + 2]; float r3 = rot[j * 9 + 3]; float r4 = rot[j * 9 + 4]; float r5 = rot[j * 9 + 5]; float r6 = rot[j * 9 + 6]; float r7 = rot[j * 9 + 7]; float r8 = rot[j * 9 + 8]; for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { float x = xyz[k * 3 + 0] - new_x; float y = xyz[k * 3 + 1] - new_y; float z = xyz[k * 3 + 2] - new_z; float x_rot = r0 * x + r3 * y + r6 * z; float y_rot = r1 * x + r4 * y + r7 * z; float z_rot = r2 * x + r5 * y + r8 * z; float d2 = y_rot * y_rot + z_rot * z_rot; if (d2 < radius2 && x_rot > hmin && x_rot < hmax) { if (cnt == 0) { for (int l = 0; l < nsample; ++l) { idx[j * nsample + l] = k; } } idx[j * nsample + cnt] = k; ++cnt; } } } } void query_cylinder_point_kernel_wrapper(int b, int n, int m, float radius, float hmin, float hmax, int nsample, const float *new_xyz, const float *xyz, const float *rot, int *idx) { cudaStream_t stream = at::cuda::getCurrentCUDAStream(); query_cylinder_point_kernel<<>>( b, n, m, radius, hmin, hmax, nsample, new_xyz, xyz, rot, idx); CUDA_CHECK_ERRORS(); } ================================================ FILE: pointnet2/_ext_src/src/group_points.cpp ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include "group_points.h" #include "utils.h" void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, const float *points, const int *idx, float *out); void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, int nsample, const float *grad_out, const int *idx, float *grad_points); at::Tensor group_points(at::Tensor points, at::Tensor idx) { CHECK_CONTIGUOUS(points); CHECK_CONTIGUOUS(idx); CHECK_IS_FLOAT(points); CHECK_IS_INT(idx); if (points.type().is_cuda()) { CHECK_CUDA(idx); } at::Tensor output = torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, at::device(points.device()).dtype(at::ScalarType::Float)); if (points.type().is_cuda()) { group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), idx.size(1), idx.size(2), points.data(), idx.data(), output.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return output; } at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { CHECK_CONTIGUOUS(grad_out); CHECK_CONTIGUOUS(idx); CHECK_IS_FLOAT(grad_out); CHECK_IS_INT(idx); if (grad_out.type().is_cuda()) { CHECK_CUDA(idx); } at::Tensor output = torch::zeros({grad_out.size(0), grad_out.size(1), n}, at::device(grad_out.device()).dtype(at::ScalarType::Float)); if (grad_out.type().is_cuda()) { group_points_grad_kernel_wrapper( grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), grad_out.data(), idx.data(), output.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return output; } ================================================ FILE: pointnet2/_ext_src/src/group_points_gpu.cu ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include #include #include "cuda_utils.h" // input: points(b, c, n) idx(b, npoints, nsample) // output: out(b, c, npoints, nsample) __global__ void group_points_kernel(int b, int c, int n, int npoints, int nsample, const float *__restrict__ points, const int *__restrict__ idx, float *__restrict__ out) { int batch_index = blockIdx.x; points += batch_index * n * c; idx += batch_index * npoints * nsample; out += batch_index * npoints * nsample * c; const int index = threadIdx.y * blockDim.x + threadIdx.x; const int stride = blockDim.y * blockDim.x; for (int i = index; i < c * npoints; i += stride) { const int l = i / npoints; const int j = i % npoints; for (int k = 0; k < nsample; ++k) { int ii = idx[j * nsample + k]; out[(l * npoints + j) * nsample + k] = points[l * n + ii]; } } } void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, const float *points, const int *idx, float *out) { cudaStream_t stream = at::cuda::getCurrentCUDAStream(); group_points_kernel<<>>( b, c, n, npoints, nsample, points, idx, out); CUDA_CHECK_ERRORS(); } // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) // output: grad_points(b, c, n) __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, int nsample, const float *__restrict__ grad_out, const int *__restrict__ idx, float *__restrict__ grad_points) { int batch_index = blockIdx.x; grad_out += batch_index * npoints * nsample * c; idx += batch_index * npoints * nsample; grad_points += batch_index * n * c; const int index = threadIdx.y * blockDim.x + threadIdx.x; const int stride = blockDim.y * blockDim.x; for (int i = index; i < c * npoints; i += stride) { const int l = i / npoints; const int j = i % npoints; for (int k = 0; k < nsample; ++k) { int ii = idx[j * nsample + k]; atomicAdd(grad_points + l * n + ii, grad_out[(l * npoints + j) * nsample + k]); } } } void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, int nsample, const float *grad_out, const int *idx, float *grad_points) { cudaStream_t stream = at::cuda::getCurrentCUDAStream(); group_points_grad_kernel<<>>( b, c, n, npoints, nsample, grad_out, idx, grad_points); CUDA_CHECK_ERRORS(); } ================================================ FILE: pointnet2/_ext_src/src/interpolate.cpp ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include "interpolate.h" #include "utils.h" void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, const float *known, float *dist2, int *idx); void three_interpolate_kernel_wrapper(int b, int c, int m, int n, const float *points, const int *idx, const float *weight, float *out); void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, const float *grad_out, const int *idx, const float *weight, float *grad_points); std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { CHECK_CONTIGUOUS(unknowns); CHECK_CONTIGUOUS(knows); CHECK_IS_FLOAT(unknowns); CHECK_IS_FLOAT(knows); if (unknowns.type().is_cuda()) { CHECK_CUDA(knows); } at::Tensor idx = torch::zeros({unknowns.size(0), unknowns.size(1), 3}, at::device(unknowns.device()).dtype(at::ScalarType::Int)); at::Tensor dist2 = torch::zeros({unknowns.size(0), unknowns.size(1), 3}, at::device(unknowns.device()).dtype(at::ScalarType::Float)); if (unknowns.type().is_cuda()) { three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), unknowns.data(), knows.data(), dist2.data(), idx.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return {dist2, idx}; } at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, at::Tensor weight) { CHECK_CONTIGUOUS(points); CHECK_CONTIGUOUS(idx); CHECK_CONTIGUOUS(weight); CHECK_IS_FLOAT(points); CHECK_IS_INT(idx); CHECK_IS_FLOAT(weight); if (points.type().is_cuda()) { CHECK_CUDA(idx); CHECK_CUDA(weight); } at::Tensor output = torch::zeros({points.size(0), points.size(1), idx.size(1)}, at::device(points.device()).dtype(at::ScalarType::Float)); if (points.type().is_cuda()) { three_interpolate_kernel_wrapper( points.size(0), points.size(1), points.size(2), idx.size(1), points.data(), idx.data(), weight.data(), output.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return output; } at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, at::Tensor weight, const int m) { CHECK_CONTIGUOUS(grad_out); CHECK_CONTIGUOUS(idx); CHECK_CONTIGUOUS(weight); CHECK_IS_FLOAT(grad_out); CHECK_IS_INT(idx); CHECK_IS_FLOAT(weight); if (grad_out.type().is_cuda()) { CHECK_CUDA(idx); CHECK_CUDA(weight); } at::Tensor output = torch::zeros({grad_out.size(0), grad_out.size(1), m}, at::device(grad_out.device()).dtype(at::ScalarType::Float)); if (grad_out.type().is_cuda()) { three_interpolate_grad_kernel_wrapper( grad_out.size(0), grad_out.size(1), grad_out.size(2), m, grad_out.data(), idx.data(), weight.data(), output.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return output; } ================================================ FILE: pointnet2/_ext_src/src/interpolate_gpu.cu ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include #include #include #include "cuda_utils.h" // input: unknown(b, n, 3) known(b, m, 3) // output: dist2(b, n, 3), idx(b, n, 3) __global__ void three_nn_kernel(int b, int n, int m, const float *__restrict__ unknown, const float *__restrict__ known, float *__restrict__ dist2, int *__restrict__ idx) { int batch_index = blockIdx.x; unknown += batch_index * n * 3; known += batch_index * m * 3; dist2 += batch_index * n * 3; idx += batch_index * n * 3; int index = threadIdx.x; int stride = blockDim.x; for (int j = index; j < n; j += stride) { float ux = unknown[j * 3 + 0]; float uy = unknown[j * 3 + 1]; float uz = unknown[j * 3 + 2]; double best1 = 1e40, best2 = 1e40, best3 = 1e40; int besti1 = 0, besti2 = 0, besti3 = 0; for (int k = 0; k < m; ++k) { float x = known[k * 3 + 0]; float y = known[k * 3 + 1]; float z = known[k * 3 + 2]; float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); if (d < best1) { best3 = best2; besti3 = besti2; best2 = best1; besti2 = besti1; best1 = d; besti1 = k; } else if (d < best2) { best3 = best2; besti3 = besti2; best2 = d; besti2 = k; } else if (d < best3) { best3 = d; besti3 = k; } } dist2[j * 3 + 0] = best1; dist2[j * 3 + 1] = best2; dist2[j * 3 + 2] = best3; idx[j * 3 + 0] = besti1; idx[j * 3 + 1] = besti2; idx[j * 3 + 2] = besti3; } } void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, const float *known, float *dist2, int *idx) { cudaStream_t stream = at::cuda::getCurrentCUDAStream(); three_nn_kernel<<>>(b, n, m, unknown, known, dist2, idx); CUDA_CHECK_ERRORS(); } // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) // output: out(b, c, n) __global__ void three_interpolate_kernel(int b, int c, int m, int n, const float *__restrict__ points, const int *__restrict__ idx, const float *__restrict__ weight, float *__restrict__ out) { int batch_index = blockIdx.x; points += batch_index * m * c; idx += batch_index * n * 3; weight += batch_index * n * 3; out += batch_index * n * c; const int index = threadIdx.y * blockDim.x + threadIdx.x; const int stride = blockDim.y * blockDim.x; for (int i = index; i < c * n; i += stride) { const int l = i / n; const int j = i % n; float w1 = weight[j * 3 + 0]; float w2 = weight[j * 3 + 1]; float w3 = weight[j * 3 + 2]; int i1 = idx[j * 3 + 0]; int i2 = idx[j * 3 + 1]; int i3 = idx[j * 3 + 2]; out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + points[l * m + i3] * w3; } } void three_interpolate_kernel_wrapper(int b, int c, int m, int n, const float *points, const int *idx, const float *weight, float *out) { cudaStream_t stream = at::cuda::getCurrentCUDAStream(); three_interpolate_kernel<<>>( b, c, m, n, points, idx, weight, out); CUDA_CHECK_ERRORS(); } // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) // output: grad_points(b, c, m) __global__ void three_interpolate_grad_kernel( int b, int c, int n, int m, const float *__restrict__ grad_out, const int *__restrict__ idx, const float *__restrict__ weight, float *__restrict__ grad_points) { int batch_index = blockIdx.x; grad_out += batch_index * n * c; idx += batch_index * n * 3; weight += batch_index * n * 3; grad_points += batch_index * m * c; const int index = threadIdx.y * blockDim.x + threadIdx.x; const int stride = blockDim.y * blockDim.x; for (int i = index; i < c * n; i += stride) { const int l = i / n; const int j = i % n; float w1 = weight[j * 3 + 0]; float w2 = weight[j * 3 + 1]; float w3 = weight[j * 3 + 2]; int i1 = idx[j * 3 + 0]; int i2 = idx[j * 3 + 1]; int i3 = idx[j * 3 + 2]; atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); } } void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, const float *grad_out, const int *idx, const float *weight, float *grad_points) { cudaStream_t stream = at::cuda::getCurrentCUDAStream(); three_interpolate_grad_kernel<<>>( b, c, n, m, grad_out, idx, weight, grad_points); CUDA_CHECK_ERRORS(); } ================================================ FILE: pointnet2/_ext_src/src/sampling.cpp ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include "sampling.h" #include "utils.h" void gather_points_kernel_wrapper(int b, int c, int n, int npoints, const float *points, const int *idx, float *out); void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, const float *grad_out, const int *idx, float *grad_points); void furthest_point_sampling_kernel_wrapper(int b, int n, int m, const float *dataset, float *temp, int *idxs); at::Tensor gather_points(at::Tensor points, at::Tensor idx) { CHECK_CONTIGUOUS(points); CHECK_CONTIGUOUS(idx); CHECK_IS_FLOAT(points); CHECK_IS_INT(idx); if (points.type().is_cuda()) { CHECK_CUDA(idx); } at::Tensor output = torch::zeros({points.size(0), points.size(1), idx.size(1)}, at::device(points.device()).dtype(at::ScalarType::Float)); if (points.type().is_cuda()) { gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), idx.size(1), points.data(), idx.data(), output.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return output; } at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { CHECK_CONTIGUOUS(grad_out); CHECK_CONTIGUOUS(idx); CHECK_IS_FLOAT(grad_out); CHECK_IS_INT(idx); if (grad_out.type().is_cuda()) { CHECK_CUDA(idx); } at::Tensor output = torch::zeros({grad_out.size(0), grad_out.size(1), n}, at::device(grad_out.device()).dtype(at::ScalarType::Float)); if (grad_out.type().is_cuda()) { gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, idx.size(1), grad_out.data(), idx.data(), output.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return output; } at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { CHECK_CONTIGUOUS(points); CHECK_IS_FLOAT(points); at::Tensor output = torch::zeros({points.size(0), nsamples}, at::device(points.device()).dtype(at::ScalarType::Int)); at::Tensor tmp = torch::full({points.size(0), points.size(1)}, 1e10, at::device(points.device()).dtype(at::ScalarType::Float)); if (points.type().is_cuda()) { furthest_point_sampling_kernel_wrapper( points.size(0), points.size(1), nsamples, points.data(), tmp.data(), output.data()); } else { TORCH_CHECK(false, "CPU not supported"); } return output; } ================================================ FILE: pointnet2/_ext_src/src/sampling_gpu.cu ================================================ // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include #include #include "cuda_utils.h" // input: points(b, c, n) idx(b, m) // output: out(b, c, m) __global__ void gather_points_kernel(int b, int c, int n, int m, const float *__restrict__ points, const int *__restrict__ idx, float *__restrict__ out) { for (int i = blockIdx.x; i < b; i += gridDim.x) { for (int l = blockIdx.y; l < c; l += gridDim.y) { for (int j = threadIdx.x; j < m; j += blockDim.x) { int a = idx[i * m + j]; out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; } } } } void gather_points_kernel_wrapper(int b, int c, int n, int npoints, const float *points, const int *idx, float *out) { gather_points_kernel<<>>(b, c, n, npoints, points, idx, out); CUDA_CHECK_ERRORS(); } // input: grad_out(b, c, m) idx(b, m) // output: grad_points(b, c, n) __global__ void gather_points_grad_kernel(int b, int c, int n, int m, const float *__restrict__ grad_out, const int *__restrict__ idx, float *__restrict__ grad_points) { for (int i = blockIdx.x; i < b; i += gridDim.x) { for (int l = blockIdx.y; l < c; l += gridDim.y) { for (int j = threadIdx.x; j < m; j += blockDim.x) { int a = idx[i * m + j]; atomicAdd(grad_points + (i * c + l) * n + a, grad_out[(i * c + l) * m + j]); } } } } void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, const float *grad_out, const int *idx, float *grad_points) { gather_points_grad_kernel<<>>( b, c, n, npoints, grad_out, idx, grad_points); CUDA_CHECK_ERRORS(); } __device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i, int idx1, int idx2) { const float v1 = dists[idx1], v2 = dists[idx2]; const int i1 = dists_i[idx1], i2 = dists_i[idx2]; dists[idx1] = max(v1, v2); dists_i[idx1] = v2 > v1 ? i2 : i1; } // Input dataset: (b, n, 3), tmp: (b, n) // Ouput idxs (b, m) template __global__ void furthest_point_sampling_kernel( int b, int n, int m, const float *__restrict__ dataset, float *__restrict__ temp, int *__restrict__ idxs) { if (m <= 0) return; __shared__ float dists[block_size]; __shared__ int dists_i[block_size]; int batch_index = blockIdx.x; dataset += batch_index * n * 3; temp += batch_index * n; idxs += batch_index * m; int tid = threadIdx.x; const int stride = block_size; int old = 0; if (threadIdx.x == 0) idxs[0] = old; __syncthreads(); for (int j = 1; j < m; j++) { int besti = 0; float best = -1; float x1 = dataset[old * 3 + 0]; float y1 = dataset[old * 3 + 1]; float z1 = dataset[old * 3 + 2]; for (int k = tid; k < n; k += stride) { float x2, y2, z2; x2 = dataset[k * 3 + 0]; y2 = dataset[k * 3 + 1]; z2 = dataset[k * 3 + 2]; float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); if (mag <= 1e-3) continue; float d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); float d2 = min(d, temp[k]); temp[k] = d2; besti = d2 > best ? k : besti; best = d2 > best ? d2 : best; } dists[tid] = best; dists_i[tid] = besti; __syncthreads(); if (block_size >= 512) { if (tid < 256) { __update(dists, dists_i, tid, tid + 256); } __syncthreads(); } if (block_size >= 256) { if (tid < 128) { __update(dists, dists_i, tid, tid + 128); } __syncthreads(); } if (block_size >= 128) { if (tid < 64) { __update(dists, dists_i, tid, tid + 64); } __syncthreads(); } if (block_size >= 64) { if (tid < 32) { __update(dists, dists_i, tid, tid + 32); } __syncthreads(); } if (block_size >= 32) { if (tid < 16) { __update(dists, dists_i, tid, tid + 16); } __syncthreads(); } if (block_size >= 16) { if (tid < 8) { __update(dists, dists_i, tid, tid + 8); } __syncthreads(); } if (block_size >= 8) { if (tid < 4) { __update(dists, dists_i, tid, tid + 4); } __syncthreads(); } if (block_size >= 4) { if (tid < 2) { __update(dists, dists_i, tid, tid + 2); } __syncthreads(); } if (block_size >= 2) { if (tid < 1) { __update(dists, dists_i, tid, tid + 1); } __syncthreads(); } old = dists_i[0]; if (tid == 0) idxs[j] = old; } } void furthest_point_sampling_kernel_wrapper(int b, int n, int m, const float *dataset, float *temp, int *idxs) { unsigned int n_threads = opt_n_threads(n); cudaStream_t stream = at::cuda::getCurrentCUDAStream(); switch (n_threads) { case 512: furthest_point_sampling_kernel<512> <<>>(b, n, m, dataset, temp, idxs); break; case 256: furthest_point_sampling_kernel<256> <<>>(b, n, m, dataset, temp, idxs); break; case 128: furthest_point_sampling_kernel<128> <<>>(b, n, m, dataset, temp, idxs); break; case 64: furthest_point_sampling_kernel<64> <<>>(b, n, m, dataset, temp, idxs); break; case 32: furthest_point_sampling_kernel<32> <<>>(b, n, m, dataset, temp, idxs); break; case 16: furthest_point_sampling_kernel<16> <<>>(b, n, m, dataset, temp, idxs); break; case 8: furthest_point_sampling_kernel<8> <<>>(b, n, m, dataset, temp, idxs); break; case 4: furthest_point_sampling_kernel<4> <<>>(b, n, m, dataset, temp, idxs); break; case 2: furthest_point_sampling_kernel<2> <<>>(b, n, m, dataset, temp, idxs); break; case 1: furthest_point_sampling_kernel<1> <<>>(b, n, m, dataset, temp, idxs); break; default: furthest_point_sampling_kernel<512> <<>>(b, n, m, dataset, temp, idxs); } CUDA_CHECK_ERRORS(); } ================================================ FILE: pointnet2/pointnet2_modules.py ================================================ # Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. ''' Pointnet2 layers. Modified based on: https://github.com/erikwijmans/Pointnet2_PyTorch Extended with the following: 1. Uniform sampling in each local region (sample_uniformly) 2. Return sampled points indices to support votenet. ''' import torch import torch.nn as nn import torch.nn.functional as F import os import sys BASE_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.append(BASE_DIR) import pointnet2_utils import pytorch_utils as pt_utils from typing import List class _PointnetSAModuleBase(nn.Module): def __init__(self): super().__init__() self.npoint = None self.groupers = None self.mlps = None def forward(self, xyz: torch.Tensor, features: torch.Tensor = None) -> (torch.Tensor, torch.Tensor): r""" Parameters ---------- xyz : torch.Tensor (B, N, 3) tensor of the xyz coordinates of the features features : torch.Tensor (B, N, C) tensor of the descriptors of the the features Returns ------- new_xyz : torch.Tensor (B, npoint, 3) tensor of the new features' xyz new_features : torch.Tensor (B, npoint, \sum_k(mlps[k][-1])) tensor of the new_features descriptors """ new_features_list = [] xyz_flipped = xyz.transpose(1, 2).contiguous() new_xyz = pointnet2_utils.gather_operation( xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint) ).transpose(1, 2).contiguous() if self.npoint is not None else None for i in range(len(self.groupers)): new_features = self.groupers[i]( xyz, new_xyz, features ) # (B, C, npoint, nsample) new_features = self.mlps[i]( new_features ) # (B, mlp[-1], npoint, nsample) new_features = F.max_pool2d( new_features, kernel_size=[1, new_features.size(3)] ) # (B, mlp[-1], npoint, 1) new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) new_features_list.append(new_features) return new_xyz, torch.cat(new_features_list, dim=1) class PointnetSAModuleMSG(_PointnetSAModuleBase): r"""Pointnet set abstrction layer with multiscale grouping Parameters ---------- npoint : int Number of features radii : list of float32 list of radii to group with nsamples : list of int32 Number of samples in each ball query mlps : list of list of int32 Spec of the pointnet before the global max_pool for each scale bn : bool Use batchnorm """ def __init__( self, *, npoint: int, radii: List[float], nsamples: List[int], mlps: List[List[int]], bn: bool = True, use_xyz: bool = True, sample_uniformly: bool = False ): super().__init__() assert len(radii) == len(nsamples) == len(mlps) self.npoint = npoint self.groupers = nn.ModuleList() self.mlps = nn.ModuleList() for i in range(len(radii)): radius = radii[i] nsample = nsamples[i] self.groupers.append( pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, sample_uniformly=sample_uniformly) if npoint is not None else pointnet2_utils.GroupAll(use_xyz) ) mlp_spec = mlps[i] if use_xyz: mlp_spec[0] += 3 self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) class PointnetSAModule(PointnetSAModuleMSG): r"""Pointnet set abstrction layer Parameters ---------- npoint : int Number of features radius : float Radius of ball nsample : int Number of samples in the ball query mlp : list Spec of the pointnet before the global max_pool bn : bool Use batchnorm """ def __init__( self, *, mlp: List[int], npoint: int = None, radius: float = None, nsample: int = None, bn: bool = True, use_xyz: bool = True ): super().__init__( mlps=[mlp], npoint=npoint, radii=[radius], nsamples=[nsample], bn=bn, use_xyz=use_xyz ) class PointnetSAModuleVotes(nn.Module): ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG with extra support for returning point indices for getting their GT votes ''' def __init__( self, *, mlp: List[int], npoint: int = None, radius: float = None, nsample: int = None, bn: bool = True, use_xyz: bool = True, pooling: str = 'max', sigma: float = None, # for RBF pooling normalize_xyz: bool = False, # noramlize local XYZ with radius sample_uniformly: bool = False, ret_unique_cnt: bool = False ): super().__init__() self.npoint = npoint self.radius = radius self.nsample = nsample self.pooling = pooling self.mlp_module = None self.use_xyz = use_xyz self.sigma = sigma if self.sigma is None: self.sigma = self.radius/2 self.normalize_xyz = normalize_xyz self.ret_unique_cnt = ret_unique_cnt if npoint is not None: self.grouper = pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, ret_grouped_xyz=True, normalize_xyz=normalize_xyz, sample_uniformly=sample_uniformly, ret_unique_cnt=ret_unique_cnt) else: self.grouper = pointnet2_utils.GroupAll(use_xyz, ret_grouped_xyz=True) mlp_spec = mlp if use_xyz and len(mlp_spec)>0: mlp_spec[0] += 3 self.mlp_module = pt_utils.SharedMLP(mlp_spec, bn=bn) def forward(self, xyz: torch.Tensor, features: torch.Tensor = None, inds: torch.Tensor = None) -> (torch.Tensor, torch.Tensor): r""" Parameters ---------- xyz : torch.Tensor (B, N, 3) tensor of the xyz coordinates of the features features : torch.Tensor (B, C, N) tensor of the descriptors of the the features inds : torch.Tensor (B, npoint) tensor that stores index to the xyz points (values in 0-N-1) Returns ------- new_xyz : torch.Tensor (B, npoint, 3) tensor of the new features' xyz new_features : torch.Tensor (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors inds: torch.Tensor (B, npoint) tensor of the inds """ xyz_flipped = xyz.transpose(1, 2).contiguous() if inds is None: inds = pointnet2_utils.furthest_point_sample(xyz, self.npoint) else: assert(inds.shape[1] == self.npoint) new_xyz = pointnet2_utils.gather_operation( xyz_flipped, inds ).transpose(1, 2).contiguous() if self.npoint is not None else None if not self.ret_unique_cnt: grouped_features, grouped_xyz = self.grouper( xyz, new_xyz, features ) # (B, C, npoint, nsample) else: grouped_features, grouped_xyz, unique_cnt = self.grouper( xyz, new_xyz, features ) # (B, C, npoint, nsample), (B,3,npoint,nsample), (B,npoint) new_features = self.mlp_module( grouped_features ) # (B, mlp[-1], npoint, nsample) if self.pooling == 'max': new_features = F.max_pool2d( new_features, kernel_size=[1, new_features.size(3)] ) # (B, mlp[-1], npoint, 1) elif self.pooling == 'avg': new_features = F.avg_pool2d( new_features, kernel_size=[1, new_features.size(3)] ) # (B, mlp[-1], npoint, 1) elif self.pooling == 'rbf': # Use radial basis function kernel for weighted sum of features (normalized by nsample and sigma) # Ref: https://en.wikipedia.org/wiki/Radial_basis_function_kernel rbf = torch.exp(-1 * grouped_xyz.pow(2).sum(1,keepdim=False) / (self.sigma**2) / 2) # (B, npoint, nsample) new_features = torch.sum(new_features * rbf.unsqueeze(1), -1, keepdim=True) / float(self.nsample) # (B, mlp[-1], npoint, 1) new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) if not self.ret_unique_cnt: return new_xyz, new_features, inds else: return new_xyz, new_features, inds, unique_cnt class PointnetSAModuleMSGVotes(nn.Module): ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG with extra support for returning point indices for getting their GT votes ''' def __init__( self, *, mlps: List[List[int]], npoint: int, radii: List[float], nsamples: List[int], bn: bool = True, use_xyz: bool = True, sample_uniformly: bool = False ): super().__init__() assert(len(mlps) == len(nsamples) == len(radii)) self.npoint = npoint self.groupers = nn.ModuleList() self.mlps = nn.ModuleList() for i in range(len(radii)): radius = radii[i] nsample = nsamples[i] self.groupers.append( pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, sample_uniformly=sample_uniformly) if npoint is not None else pointnet2_utils.GroupAll(use_xyz) ) mlp_spec = mlps[i] if use_xyz: mlp_spec[0] += 3 self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) def forward(self, xyz: torch.Tensor, features: torch.Tensor = None, inds: torch.Tensor = None) -> (torch.Tensor, torch.Tensor): r""" Parameters ---------- xyz : torch.Tensor (B, N, 3) tensor of the xyz coordinates of the features features : torch.Tensor (B, C, C) tensor of the descriptors of the the features inds : torch.Tensor (B, npoint) tensor that stores index to the xyz points (values in 0-N-1) Returns ------- new_xyz : torch.Tensor (B, npoint, 3) tensor of the new features' xyz new_features : torch.Tensor (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors inds: torch.Tensor (B, npoint) tensor of the inds """ new_features_list = [] xyz_flipped = xyz.transpose(1, 2).contiguous() if inds is None: inds = pointnet2_utils.furthest_point_sample(xyz, self.npoint) new_xyz = pointnet2_utils.gather_operation( xyz_flipped, inds ).transpose(1, 2).contiguous() if self.npoint is not None else None for i in range(len(self.groupers)): new_features = self.groupers[i]( xyz, new_xyz, features ) # (B, C, npoint, nsample) new_features = self.mlps[i]( new_features ) # (B, mlp[-1], npoint, nsample) new_features = F.max_pool2d( new_features, kernel_size=[1, new_features.size(3)] ) # (B, mlp[-1], npoint, 1) new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) new_features_list.append(new_features) return new_xyz, torch.cat(new_features_list, dim=1), inds class PointnetFPModule(nn.Module): r"""Propigates the features of one set to another Parameters ---------- mlp : list Pointnet module parameters bn : bool Use batchnorm """ def __init__(self, *, mlp: List[int], bn: bool = True): super().__init__() self.mlp = pt_utils.SharedMLP(mlp, bn=bn) def forward( self, unknown: torch.Tensor, known: torch.Tensor, unknow_feats: torch.Tensor, known_feats: torch.Tensor ) -> torch.Tensor: r""" Parameters ---------- unknown : torch.Tensor (B, n, 3) tensor of the xyz positions of the unknown features known : torch.Tensor (B, m, 3) tensor of the xyz positions of the known features unknow_feats : torch.Tensor (B, C1, n) tensor of the features to be propigated to known_feats : torch.Tensor (B, C2, m) tensor of features to be propigated Returns ------- new_features : torch.Tensor (B, mlp[-1], n) tensor of the features of the unknown features """ if known is not None: dist, idx = pointnet2_utils.three_nn(unknown, known) dist_recip = 1.0 / (dist + 1e-8) norm = torch.sum(dist_recip, dim=2, keepdim=True) weight = dist_recip / norm interpolated_feats = pointnet2_utils.three_interpolate( known_feats, idx, weight ) else: interpolated_feats = known_feats.expand( *known_feats.size()[0:2], unknown.size(1) ) if unknow_feats is not None: new_features = torch.cat([interpolated_feats, unknow_feats], dim=1) #(B, C2 + C1, n) else: new_features = interpolated_feats new_features = new_features.unsqueeze(-1) new_features = self.mlp(new_features) return new_features.squeeze(-1) class PointnetLFPModuleMSG(nn.Module): ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG learnable feature propagation layer.''' def __init__( self, *, mlps: List[List[int]], radii: List[float], nsamples: List[int], post_mlp: List[int], bn: bool = True, use_xyz: bool = True, sample_uniformly: bool = False ): super().__init__() assert(len(mlps) == len(nsamples) == len(radii)) self.post_mlp = pt_utils.SharedMLP(post_mlp, bn=bn) self.groupers = nn.ModuleList() self.mlps = nn.ModuleList() for i in range(len(radii)): radius = radii[i] nsample = nsamples[i] self.groupers.append( pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, sample_uniformly=sample_uniformly) ) mlp_spec = mlps[i] if use_xyz: mlp_spec[0] += 3 self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) def forward(self, xyz2: torch.Tensor, xyz1: torch.Tensor, features2: torch.Tensor, features1: torch.Tensor) -> torch.Tensor: r""" Propagate features from xyz1 to xyz2. Parameters ---------- xyz2 : torch.Tensor (B, N2, 3) tensor of the xyz coordinates of the features xyz1 : torch.Tensor (B, N1, 3) tensor of the xyz coordinates of the features features2 : torch.Tensor (B, C2, N2) tensor of the descriptors of the the features features1 : torch.Tensor (B, C1, N1) tensor of the descriptors of the the features Returns ------- new_features1 : torch.Tensor (B, \sum_k(mlps[k][-1]), N1) tensor of the new_features descriptors """ new_features_list = [] for i in range(len(self.groupers)): new_features = self.groupers[i]( xyz1, xyz2, features1 ) # (B, C1, N2, nsample) new_features = self.mlps[i]( new_features ) # (B, mlp[-1], N2, nsample) new_features = F.max_pool2d( new_features, kernel_size=[1, new_features.size(3)] ) # (B, mlp[-1], N2, 1) new_features = new_features.squeeze(-1) # (B, mlp[-1], N2) if features2 is not None: new_features = torch.cat([new_features, features2], dim=1) #(B, mlp[-1] + C2, N2) new_features = new_features.unsqueeze(-1) new_features = self.post_mlp(new_features) new_features_list.append(new_features) return torch.cat(new_features_list, dim=1).squeeze(-1) if __name__ == "__main__": from torch.autograd import Variable torch.manual_seed(1) torch.cuda.manual_seed_all(1) xyz = Variable(torch.randn(2, 9, 3).cuda(), requires_grad=True) xyz_feats = Variable(torch.randn(2, 9, 6).cuda(), requires_grad=True) test_module = PointnetSAModuleMSG( npoint=2, radii=[5.0, 10.0], nsamples=[6, 3], mlps=[[9, 3], [9, 6]] ) test_module.cuda() print(test_module(xyz, xyz_feats)) for _ in range(1): _, new_features = test_module(xyz, xyz_feats) new_features.backward( torch.cuda.FloatTensor(*new_features.size()).fill_(1) ) print(new_features) print(xyz.grad) ================================================ FILE: pointnet2/pointnet2_utils.py ================================================ # Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. ''' Modified based on: https://github.com/erikwijmans/Pointnet2_PyTorch ''' from __future__ import ( division, absolute_import, with_statement, print_function, unicode_literals, ) import torch from torch.autograd import Function import torch.nn as nn import pytorch_utils as pt_utils import sys try: import builtins except: import __builtin__ as builtins try: import pointnet2._ext as _ext except ImportError: if not getattr(builtins, "__POINTNET2_SETUP__", False): raise ImportError( "Could not import _ext module.\n" "Please see the setup instructions in the README: " "https://github.com/erikwijmans/Pointnet2_PyTorch/blob/master/README.rst" ) if False: # Workaround for type hints without depending on the `typing` module from typing import * class RandomDropout(nn.Module): def __init__(self, p=0.5, inplace=False): super(RandomDropout, self).__init__() self.p = p self.inplace = inplace def forward(self, X): theta = torch.Tensor(1).uniform_(0, self.p)[0] return pt_utils.feature_dropout_no_scaling(X, theta, self.train, self.inplace) class FurthestPointSampling(Function): @staticmethod def forward(ctx, xyz, npoint): # type: (Any, torch.Tensor, int) -> torch.Tensor r""" Uses iterative furthest point sampling to select a set of npoint features that have the largest minimum distance Parameters ---------- xyz : torch.Tensor (B, N, 3) tensor where N > npoint npoint : int32 number of features in the sampled set Returns ------- torch.Tensor (B, npoint) tensor containing the set """ return _ext.furthest_point_sampling(xyz, npoint) @staticmethod def backward(xyz, a=None): return None, None furthest_point_sample = FurthestPointSampling.apply class GatherOperation(Function): @staticmethod def forward(ctx, features, idx): # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor r""" Parameters ---------- features : torch.Tensor (B, C, N) tensor idx : torch.Tensor (B, npoint) tensor of the features to gather Returns ------- torch.Tensor (B, C, npoint) tensor """ _, C, N = features.size() ctx.for_backwards = (idx, C, N) return _ext.gather_points(features, idx) @staticmethod def backward(ctx, grad_out): idx, C, N = ctx.for_backwards grad_features = _ext.gather_points_grad(grad_out.contiguous(), idx, N) return grad_features, None gather_operation = GatherOperation.apply class ThreeNN(Function): @staticmethod def forward(ctx, unknown, known): # type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] r""" Find the three nearest neighbors of unknown in known Parameters ---------- unknown : torch.Tensor (B, n, 3) tensor of known features known : torch.Tensor (B, m, 3) tensor of unknown features Returns ------- dist : torch.Tensor (B, n, 3) l2 distance to the three nearest neighbors idx : torch.Tensor (B, n, 3) index of 3 nearest neighbors """ dist2, idx = _ext.three_nn(unknown, known) return torch.sqrt(dist2), idx @staticmethod def backward(ctx, a=None, b=None): return None, None three_nn = ThreeNN.apply class ThreeInterpolate(Function): @staticmethod def forward(ctx, features, idx, weight): # type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor r""" Performs weight linear interpolation on 3 features Parameters ---------- features : torch.Tensor (B, c, m) Features descriptors to be interpolated from idx : torch.Tensor (B, n, 3) three nearest neighbors of the target features in features weight : torch.Tensor (B, n, 3) weights Returns ------- torch.Tensor (B, c, n) tensor of the interpolated features """ B, c, m = features.size() n = idx.size(1) ctx.three_interpolate_for_backward = (idx, weight, m) return _ext.three_interpolate(features, idx, weight) @staticmethod def backward(ctx, grad_out): # type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor] r""" Parameters ---------- grad_out : torch.Tensor (B, c, n) tensor with gradients of ouputs Returns ------- grad_features : torch.Tensor (B, c, m) tensor with gradients of features None None """ idx, weight, m = ctx.three_interpolate_for_backward grad_features = _ext.three_interpolate_grad( grad_out.contiguous(), idx, weight, m ) return grad_features, None, None three_interpolate = ThreeInterpolate.apply class GroupingOperation(Function): @staticmethod def forward(ctx, features, idx): # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor r""" Parameters ---------- features : torch.Tensor (B, C, N) tensor of features to group idx : torch.Tensor (B, npoint, nsample) tensor containing the indicies of features to group with Returns ------- torch.Tensor (B, C, npoint, nsample) tensor """ B, nfeatures, nsample = idx.size() _, C, N = features.size() ctx.for_backwards = (idx, N) return _ext.group_points(features, idx) @staticmethod def backward(ctx, grad_out): # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor] r""" Parameters ---------- grad_out : torch.Tensor (B, C, npoint, nsample) tensor of the gradients of the output from forward Returns ------- torch.Tensor (B, C, N) gradient of the features None """ idx, N = ctx.for_backwards grad_features = _ext.group_points_grad(grad_out.contiguous(), idx, N) return grad_features, None grouping_operation = GroupingOperation.apply class BallQuery(Function): @staticmethod def forward(ctx, radius, nsample, xyz, new_xyz): # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor r""" Parameters ---------- radius : float radius of the balls nsample : int maximum number of features in the balls xyz : torch.Tensor (B, N, 3) xyz coordinates of the features new_xyz : torch.Tensor (B, npoint, 3) centers of the ball query Returns ------- torch.Tensor (B, npoint, nsample) tensor with the indicies of the features that form the query balls """ return _ext.ball_query(new_xyz, xyz, radius, nsample) @staticmethod def backward(ctx, a=None): return None, None, None, None ball_query = BallQuery.apply class QueryAndGroup(nn.Module): r""" Groups with a ball query of radius Parameters --------- radius : float32 Radius of ball nsample : int32 Maximum number of features to gather in the ball """ def __init__(self, radius, nsample, use_xyz=True, ret_grouped_xyz=False, normalize_xyz=False, sample_uniformly=False, ret_unique_cnt=False): # type: (QueryAndGroup, float, int, bool) -> None super(QueryAndGroup, self).__init__() self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz self.ret_grouped_xyz = ret_grouped_xyz self.normalize_xyz = normalize_xyz self.sample_uniformly = sample_uniformly self.ret_unique_cnt = ret_unique_cnt if self.ret_unique_cnt: assert(self.sample_uniformly) def forward(self, xyz, new_xyz, features=None): # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor] r""" Parameters ---------- xyz : torch.Tensor xyz coordinates of the features (B, N, 3) new_xyz : torch.Tensor centriods (B, npoint, 3) features : torch.Tensor Descriptors of the features (B, C, N) Returns ------- new_features : torch.Tensor (B, 3 + C, npoint, nsample) tensor """ idx = ball_query(self.radius, self.nsample, xyz, new_xyz) if self.sample_uniformly: unique_cnt = torch.zeros((idx.shape[0], idx.shape[1])) for i_batch in range(idx.shape[0]): for i_region in range(idx.shape[1]): unique_ind = torch.unique(idx[i_batch, i_region, :]) num_unique = unique_ind.shape[0] unique_cnt[i_batch, i_region] = num_unique sample_ind = torch.randint(0, num_unique, (self.nsample - num_unique,), dtype=torch.long) all_ind = torch.cat((unique_ind, unique_ind[sample_ind])) idx[i_batch, i_region, :] = all_ind xyz_trans = xyz.transpose(1, 2).contiguous() grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1) if self.normalize_xyz: grouped_xyz /= self.radius if features is not None: grouped_features = grouping_operation(features, idx) if self.use_xyz: new_features = torch.cat( [grouped_xyz, grouped_features], dim=1 ) # (B, C + 3, npoint, nsample) else: new_features = grouped_features else: assert ( self.use_xyz ), "Cannot have not features and not use xyz as a feature!" new_features = grouped_xyz ret = [new_features] if self.ret_grouped_xyz: ret.append(grouped_xyz) if self.ret_unique_cnt: ret.append(unique_cnt) if len(ret) == 1: return ret[0] else: return tuple(ret) class GroupAll(nn.Module): r""" Groups all features Parameters --------- """ def __init__(self, use_xyz=True, ret_grouped_xyz=False): # type: (GroupAll, bool) -> None super(GroupAll, self).__init__() self.use_xyz = use_xyz def forward(self, xyz, new_xyz, features=None): # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor] r""" Parameters ---------- xyz : torch.Tensor xyz coordinates of the features (B, N, 3) new_xyz : torch.Tensor Ignored features : torch.Tensor Descriptors of the features (B, C, N) Returns ------- new_features : torch.Tensor (B, C + 3, 1, N) tensor """ grouped_xyz = xyz.transpose(1, 2).unsqueeze(2) if features is not None: grouped_features = features.unsqueeze(2) if self.use_xyz: new_features = torch.cat( [grouped_xyz, grouped_features], dim=1 ) # (B, 3 + C, 1, N) else: new_features = grouped_features else: new_features = grouped_xyz if self.ret_grouped_xyz: return new_features, grouped_xyz else: return new_features class CylinderQuery(Function): @staticmethod def forward(ctx, radius, hmin, hmax, nsample, xyz, new_xyz, rot): # type: (Any, float, float, float, int, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor r""" Parameters ---------- radius : float radius of the cylinders hmin, hmax : float endpoints of cylinder height in x-rotation axis nsample : int maximum number of features in the cylinders xyz : torch.Tensor (B, N, 3) xyz coordinates of the features new_xyz : torch.Tensor (B, npoint, 3) centers of the cylinder query rot: torch.Tensor (B, npoint, 9) flatten rotation matrices from cylinder frame to world frame Returns ------- torch.Tensor (B, npoint, nsample) tensor with the indicies of the features that form the query balls """ return _ext.cylinder_query(new_xyz, xyz, rot, radius, hmin, hmax, nsample) @staticmethod def backward(ctx, a=None): return None, None, None, None, None, None, None cylinder_query = CylinderQuery.apply class CylinderQueryAndGroup(nn.Module): r""" Groups with a cylinder query of radius and height Parameters --------- radius : float32 Radius of cylinder hmin, hmax: float32 endpoints of cylinder height in x-rotation axis nsample : int32 Maximum number of features to gather in the ball """ def __init__(self, radius, hmin, hmax, nsample, use_xyz=True, ret_grouped_xyz=False, normalize_xyz=False, rotate_xyz=True, sample_uniformly=False, ret_unique_cnt=False): super(CylinderQueryAndGroup, self).__init__() self.radius, self.nsample, self.hmin, self.hmax, = radius, nsample, hmin, hmax self.use_xyz = use_xyz self.ret_grouped_xyz = ret_grouped_xyz self.normalize_xyz = normalize_xyz self.rotate_xyz = rotate_xyz self.sample_uniformly = sample_uniformly self.ret_unique_cnt = ret_unique_cnt if self.ret_unique_cnt: assert(self.sample_uniformly) def forward(self, xyz, new_xyz, rot, features=None): r""" Parameters ---------- xyz : torch.Tensor xyz coordinates of the features (B, N, 3) new_xyz : torch.Tensor centriods (B, npoint, 3) rot : torch.Tensor rotation matrices (B, npoint, 3, 3) features : torch.Tensor Descriptors of the features (B, C, N) Returns ------- new_features : torch.Tensor (B, 3 + C, npoint, nsample) tensor """ B, npoint, _ = new_xyz.size() idx = cylinder_query(self.radius, self.hmin, self.hmax, self.nsample, xyz, new_xyz, rot.view(B, npoint, 9)) if self.sample_uniformly: unique_cnt = torch.zeros((idx.shape[0], idx.shape[1])) for i_batch in range(idx.shape[0]): for i_region in range(idx.shape[1]): unique_ind = torch.unique(idx[i_batch, i_region, :]) num_unique = unique_ind.shape[0] unique_cnt[i_batch, i_region] = num_unique sample_ind = torch.randint(0, num_unique, (self.nsample - num_unique,), dtype=torch.long) all_ind = torch.cat((unique_ind, unique_ind[sample_ind])) idx[i_batch, i_region, :] = all_ind xyz_trans = xyz.transpose(1, 2).contiguous() grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1) if self.normalize_xyz: grouped_xyz /= self.radius if self.rotate_xyz: grouped_xyz_ = grouped_xyz.permute(0, 2, 3, 1).contiguous() # (B, npoint, nsample, 3) grouped_xyz_ = torch.matmul(grouped_xyz_, rot) grouped_xyz = grouped_xyz_.permute(0, 3, 1, 2).contiguous() if features is not None: grouped_features = grouping_operation(features, idx) if self.use_xyz: new_features = torch.cat( [grouped_xyz, grouped_features], dim=1 ) # (B, C + 3, npoint, nsample) else: new_features = grouped_features else: assert ( self.use_xyz ), "Cannot have not features and not use xyz as a feature!" new_features = grouped_xyz ret = [new_features] if self.ret_grouped_xyz: ret.append(grouped_xyz) if self.ret_unique_cnt: ret.append(unique_cnt) if len(ret) == 1: return ret[0] else: return tuple(ret) ================================================ FILE: pointnet2/pytorch_utils.py ================================================ # Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. ''' Modified based on Ref: https://github.com/erikwijmans/Pointnet2_PyTorch ''' import torch import torch.nn as nn from typing import List, Tuple class SharedMLP(nn.Sequential): def __init__( self, args: List[int], *, bn: bool = False, activation=nn.ReLU(inplace=True), preact: bool = False, first: bool = False, name: str = "" ): super().__init__() for i in range(len(args) - 1): self.add_module( name + 'layer{}'.format(i), Conv2d( args[i], args[i + 1], bn=(not first or not preact or (i != 0)) and bn, activation=activation if (not first or not preact or (i != 0)) else None, preact=preact ) ) class _BNBase(nn.Sequential): def __init__(self, in_size, batch_norm=None, name=""): super().__init__() self.add_module(name + "bn", batch_norm(in_size)) nn.init.constant_(self[0].weight, 1.0) nn.init.constant_(self[0].bias, 0) class BatchNorm1d(_BNBase): def __init__(self, in_size: int, *, name: str = ""): super().__init__(in_size, batch_norm=nn.BatchNorm1d, name=name) class BatchNorm2d(_BNBase): def __init__(self, in_size: int, name: str = ""): super().__init__(in_size, batch_norm=nn.BatchNorm2d, name=name) class BatchNorm3d(_BNBase): def __init__(self, in_size: int, name: str = ""): super().__init__(in_size, batch_norm=nn.BatchNorm3d, name=name) class _ConvBase(nn.Sequential): def __init__( self, in_size, out_size, kernel_size, stride, padding, activation, bn, init, conv=None, batch_norm=None, bias=True, preact=False, name="" ): super().__init__() bias = bias and (not bn) conv_unit = conv( in_size, out_size, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias ) init(conv_unit.weight) if bias: nn.init.constant_(conv_unit.bias, 0) if bn: if not preact: bn_unit = batch_norm(out_size) else: bn_unit = batch_norm(in_size) if preact: if bn: self.add_module(name + 'bn', bn_unit) if activation is not None: self.add_module(name + 'activation', activation) self.add_module(name + 'conv', conv_unit) if not preact: if bn: self.add_module(name + 'bn', bn_unit) if activation is not None: self.add_module(name + 'activation', activation) class Conv1d(_ConvBase): def __init__( self, in_size: int, out_size: int, *, kernel_size: int = 1, stride: int = 1, padding: int = 0, activation=nn.ReLU(inplace=True), bn: bool = False, init=nn.init.kaiming_normal_, bias: bool = True, preact: bool = False, name: str = "" ): super().__init__( in_size, out_size, kernel_size, stride, padding, activation, bn, init, conv=nn.Conv1d, batch_norm=BatchNorm1d, bias=bias, preact=preact, name=name ) class Conv2d(_ConvBase): def __init__( self, in_size: int, out_size: int, *, kernel_size: Tuple[int, int] = (1, 1), stride: Tuple[int, int] = (1, 1), padding: Tuple[int, int] = (0, 0), activation=nn.ReLU(inplace=True), bn: bool = False, init=nn.init.kaiming_normal_, bias: bool = True, preact: bool = False, name: str = "" ): super().__init__( in_size, out_size, kernel_size, stride, padding, activation, bn, init, conv=nn.Conv2d, batch_norm=BatchNorm2d, bias=bias, preact=preact, name=name ) class Conv3d(_ConvBase): def __init__( self, in_size: int, out_size: int, *, kernel_size: Tuple[int, int, int] = (1, 1, 1), stride: Tuple[int, int, int] = (1, 1, 1), padding: Tuple[int, int, int] = (0, 0, 0), activation=nn.ReLU(inplace=True), bn: bool = False, init=nn.init.kaiming_normal_, bias: bool = True, preact: bool = False, name: str = "" ): super().__init__( in_size, out_size, kernel_size, stride, padding, activation, bn, init, conv=nn.Conv3d, batch_norm=BatchNorm3d, bias=bias, preact=preact, name=name ) class FC(nn.Sequential): def __init__( self, in_size: int, out_size: int, *, activation=nn.ReLU(inplace=True), bn: bool = False, init=None, preact: bool = False, name: str = "" ): super().__init__() fc = nn.Linear(in_size, out_size, bias=not bn) if init is not None: init(fc.weight) if not bn: nn.init.constant_(fc.bias, 0) if preact: if bn: self.add_module(name + 'bn', BatchNorm1d(in_size)) if activation is not None: self.add_module(name + 'activation', activation) self.add_module(name + 'fc', fc) if not preact: if bn: self.add_module(name + 'bn', BatchNorm1d(out_size)) if activation is not None: self.add_module(name + 'activation', activation) def set_bn_momentum_default(bn_momentum): def fn(m): if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)): m.momentum = bn_momentum return fn class BNMomentumScheduler(object): def __init__( self, model, bn_lambda, last_epoch=-1, setter=set_bn_momentum_default ): if not isinstance(model, nn.Module): raise RuntimeError( "Class '{}' is not a PyTorch nn Module".format( type(model).__name__ ) ) self.model = model self.setter = setter self.lmbd = bn_lambda self.step(last_epoch + 1) self.last_epoch = last_epoch def step(self, epoch=None): if epoch is None: epoch = self.last_epoch + 1 self.last_epoch = epoch self.model.apply(self.setter(self.lmbd(epoch))) ================================================ FILE: pointnet2/setup.py ================================================ # Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from setuptools import setup from torch.utils.cpp_extension import BuildExtension, CUDAExtension import glob import os ROOT = os.path.dirname(os.path.abspath(__file__)) _ext_src_root = "_ext_src" _ext_sources = glob.glob("{}/src/*.cpp".format(_ext_src_root)) + glob.glob( "{}/src/*.cu".format(_ext_src_root) ) _ext_headers = glob.glob("{}/include/*".format(_ext_src_root)) setup( name='pointnet2', ext_modules=[ CUDAExtension( name='pointnet2._ext', sources=_ext_sources, extra_compile_args={ "cxx": ["-O2", "-I{}".format("{}/{}/include".format(ROOT, _ext_src_root))], "nvcc": ["-O2", "-I{}".format("{}/{}/include".format(ROOT, _ext_src_root))], }, ) ], cmdclass={ 'build_ext': BuildExtension } ) ================================================ FILE: requirements.txt ================================================ torch>=1.8 tensorboard==2.3 numpy scipy open3d>=0.8 Pillow tqdm MinkowskiEngine==0.5.4 ================================================ FILE: test.py ================================================ import os import sys import numpy as np import argparse import time import torch from torch.utils.data import DataLoader from graspnetAPI.graspnet_eval import GraspGroup, GraspNetEval ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.append(os.path.join(ROOT_DIR, 'pointnet2')) sys.path.append(os.path.join(ROOT_DIR, 'utils')) sys.path.append(os.path.join(ROOT_DIR, 'models')) sys.path.append(os.path.join(ROOT_DIR, 'dataset')) from models.graspnet import GraspNet, pred_decode from dataset.graspnet_dataset import GraspNetDataset, minkowski_collate_fn from utils.collision_detector import ModelFreeCollisionDetector parser = argparse.ArgumentParser() parser.add_argument('--dataset_root', default=None, required=True) parser.add_argument('--checkpoint_path', help='Model checkpoint path', default=None, required=True) parser.add_argument('--dump_dir', help='Dump dir to save outputs', default=None, required=True) parser.add_argument('--seed_feat_dim', default=512, type=int, help='Point wise feature dim') parser.add_argument('--camera', default='kinect', help='Camera split [realsense/kinect]') parser.add_argument('--num_point', type=int, default=15000, help='Point Number [default: 15000]') parser.add_argument('--batch_size', type=int, default=1, help='Batch Size during inference [default: 1]') parser.add_argument('--voxel_size', type=float, default=0.005, help='Voxel Size for sparse convolution') parser.add_argument('--collision_thresh', type=float, default=0.01, help='Collision Threshold in collision detection [default: 0.01]') parser.add_argument('--voxel_size_cd', type=float, default=0.01, help='Voxel Size for collision detection') parser.add_argument('--infer', action='store_true', default=False) parser.add_argument('--eval', action='store_true', default=False) cfgs = parser.parse_args() # ------------------------------------------------------------------------- GLOBAL CONFIG BEG if not os.path.exists(cfgs.dump_dir): os.mkdir(cfgs.dump_dir) # Init datasets and dataloaders def my_worker_init_fn(worker_id): np.random.seed(np.random.get_state()[1][0] + worker_id) pass def inference(): test_dataset = GraspNetDataset(cfgs.dataset_root, split='test_seen', camera=cfgs.camera, num_points=cfgs.num_point, voxel_size=cfgs.voxel_size, remove_outlier=True, augment=False, load_label=False) print('Test dataset length: ', len(test_dataset)) scene_list = test_dataset.scene_list() test_dataloader = DataLoader(test_dataset, batch_size=cfgs.batch_size, shuffle=False, num_workers=0, worker_init_fn=my_worker_init_fn, collate_fn=minkowski_collate_fn) print('Test dataloader length: ', len(test_dataloader)) # Init the model net = GraspNet(seed_feat_dim=cfgs.seed_feat_dim, is_training=False) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device) # Load checkpoint checkpoint = torch.load(cfgs.checkpoint_path) net.load_state_dict(checkpoint['model_state_dict']) start_epoch = checkpoint['epoch'] print("-> loaded checkpoint %s (epoch: %d)" % (cfgs.checkpoint_path, start_epoch)) batch_interval = 100 net.eval() tic = time.time() for batch_idx, batch_data in enumerate(test_dataloader): for key in batch_data: if 'list' in key: for i in range(len(batch_data[key])): for j in range(len(batch_data[key][i])): batch_data[key][i][j] = batch_data[key][i][j].to(device) else: batch_data[key] = batch_data[key].to(device) # Forward pass with torch.no_grad(): end_points = net(batch_data) grasp_preds = pred_decode(end_points) # Dump results for evaluation for i in range(cfgs.batch_size): data_idx = batch_idx * cfgs.batch_size + i preds = grasp_preds[i].detach().cpu().numpy() gg = GraspGroup(preds) # collision detection if cfgs.collision_thresh > 0: cloud = test_dataset.get_data(data_idx, return_raw_cloud=True) mfcdetector = ModelFreeCollisionDetector(cloud, voxel_size=cfgs.voxel_size_cd) collision_mask = mfcdetector.detect(gg, approach_dist=0.05, collision_thresh=cfgs.collision_thresh) gg = gg[~collision_mask] # save grasps save_dir = os.path.join(cfgs.dump_dir, scene_list[data_idx], cfgs.camera) save_path = os.path.join(save_dir, str(data_idx % 256).zfill(4) + '.npy') if not os.path.exists(save_dir): os.makedirs(save_dir) gg.save_npy(save_path) if (batch_idx + 1) % batch_interval == 0: toc = time.time() print('Eval batch: %d, time: %fs' % (batch_idx + 1, (toc - tic) / batch_interval)) tic = time.time() def evaluate(dump_dir): ge = GraspNetEval(root=cfgs.dataset_root, camera=cfgs.camera, split='test_seen') res, ap = ge.eval_seen(dump_folder=dump_dir, proc=6) save_dir = os.path.join(cfgs.dump_dir, 'ap_{}.npy'.format(cfgs.camera)) np.save(save_dir, res) if __name__ == '__main__': if cfgs.infer: inference() if cfgs.eval: evaluate(cfgs.dump_dir) ================================================ FILE: train.py ================================================ import os import sys import numpy as np from datetime import datetime import argparse import torch import torch.optim as optim from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.append(os.path.join(ROOT_DIR, 'pointnet2')) sys.path.append(os.path.join(ROOT_DIR, 'utils')) sys.path.append(os.path.join(ROOT_DIR, 'models')) sys.path.append(os.path.join(ROOT_DIR, 'dataset')) from models.graspnet import GraspNet from models.loss import get_loss from dataset.graspnet_dataset import GraspNetDataset, minkowski_collate_fn, load_grasp_labels parser = argparse.ArgumentParser() parser.add_argument('--dataset_root', default=None, required=True) parser.add_argument('--camera', default='kinect', help='Camera split [realsense/kinect]') parser.add_argument('--checkpoint_path', help='Model checkpoint path', default=None) parser.add_argument('--model_name', type=str, default=None) parser.add_argument('--log_dir', default='logs/log') parser.add_argument('--num_point', type=int, default=15000, help='Point Number [default: 20000]') parser.add_argument('--seed_feat_dim', default=512, type=int, help='Point wise feature dim') parser.add_argument('--voxel_size', type=float, default=0.005, help='Voxel Size to process point clouds ') parser.add_argument('--max_epoch', type=int, default=10, help='Epoch to run [default: 18]') parser.add_argument('--batch_size', type=int, default=4, help='Batch Size during training [default: 2]') parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate [default: 0.001]') parser.add_argument('--resume', action='store_true', default=False, help='Whether to resume from checkpoint') cfgs = parser.parse_args() # ------------------------------------------------------------------------- GLOBAL CONFIG BEG EPOCH_CNT = 0 CHECKPOINT_PATH = cfgs.checkpoint_path if cfgs.checkpoint_path is not None and cfgs.resume else None if not os.path.exists(cfgs.log_dir): os.makedirs(cfgs.log_dir) LOG_FOUT = open(os.path.join(cfgs.log_dir, 'log_train.txt'), 'a') LOG_FOUT.write(str(cfgs) + '\n') def log_string(out_str): LOG_FOUT.write(out_str + '\n') LOG_FOUT.flush() print(out_str) # Init datasets and dataloaders def my_worker_init_fn(worker_id): np.random.seed(np.random.get_state()[1][0] + worker_id) pass grasp_labels = load_grasp_labels(cfgs.dataset_root) TRAIN_DATASET = GraspNetDataset(cfgs.dataset_root, grasp_labels=grasp_labels, camera=cfgs.camera, split='train', num_points=cfgs.num_point, voxel_size=cfgs.voxel_size, remove_outlier=True, augment=True, load_label=True) print('train dataset length: ', len(TRAIN_DATASET)) TRAIN_DATALOADER = DataLoader(TRAIN_DATASET, batch_size=cfgs.batch_size, shuffle=True, num_workers=0, worker_init_fn=my_worker_init_fn, collate_fn=minkowski_collate_fn) print('train dataloader length: ', len(TRAIN_DATALOADER)) net = GraspNet(seed_feat_dim=cfgs.seed_feat_dim, is_training=True) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device) # Load the Adam optimizer optimizer = optim.Adam(net.parameters(), lr=cfgs.learning_rate) start_epoch = 0 if CHECKPOINT_PATH is not None and os.path.isfile(CHECKPOINT_PATH): checkpoint = torch.load(CHECKPOINT_PATH) net.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) start_epoch = checkpoint['epoch'] log_string("-> loaded checkpoint %s (epoch: %d)" % (CHECKPOINT_PATH, start_epoch)) # TensorBoard Visualizers TRAIN_WRITER = SummaryWriter(os.path.join(cfgs.log_dir, 'train')) def get_current_lr(epoch): lr = cfgs.learning_rate lr = lr * (0.95 ** epoch) return lr def adjust_learning_rate(optimizer, epoch): lr = get_current_lr(epoch) for param_group in optimizer.param_groups: param_group['lr'] = lr def train_one_epoch(): stat_dict = {} # collect statistics adjust_learning_rate(optimizer, EPOCH_CNT) net.train() batch_interval = 20 for batch_idx, batch_data_label in enumerate(TRAIN_DATALOADER): for key in batch_data_label: if 'list' in key: for i in range(len(batch_data_label[key])): for j in range(len(batch_data_label[key][i])): batch_data_label[key][i][j] = batch_data_label[key][i][j].to(device) else: batch_data_label[key] = batch_data_label[key].to(device) end_points = net(batch_data_label) loss, end_points = get_loss(end_points) loss.backward() optimizer.step() optimizer.zero_grad() for key in end_points: if 'loss' in key or 'acc' in key or 'prec' in key or 'recall' in key or 'count' in key: if key not in stat_dict: stat_dict[key] = 0 stat_dict[key] += end_points[key].item() if (batch_idx + 1) % batch_interval == 0: log_string(' ----epoch: %03d ---- batch: %03d ----' % (EPOCH_CNT, batch_idx + 1)) for key in sorted(stat_dict.keys()): TRAIN_WRITER.add_scalar(key, stat_dict[key] / batch_interval, (EPOCH_CNT * len(TRAIN_DATALOADER) + batch_idx) * cfgs.batch_size) log_string('mean %s: %f' % (key, stat_dict[key] / batch_interval)) stat_dict[key] = 0 def train(start_epoch): global EPOCH_CNT for epoch in range(start_epoch, cfgs.max_epoch): EPOCH_CNT = epoch log_string('**** EPOCH %03d ****' % epoch) log_string('Current learning rate: %f' % (get_current_lr(epoch))) log_string(str(datetime.now())) # Reset numpy seed. # REF: https://github.com/pytorch/pytorch/issues/5059 np.random.seed() train_one_epoch() save_dict = {'epoch': epoch + 1, 'optimizer_state_dict': optimizer.state_dict(), 'model_state_dict': net.state_dict()} torch.save(save_dict, os.path.join(cfgs.log_dir, cfgs.model_name + '_epoch' + str(epoch + 1).zfill(2) + '.tar')) if __name__ == '__main__': train(start_epoch) ================================================ FILE: utils/collision_detector.py ================================================ """ Collision detection to remove collided grasp pose predictions. Author: chenxi-wang """ import os import sys import numpy as np import open3d as o3d class ModelFreeCollisionDetector(): """ Collision detection in scenes without object labels. Current finger width and length are fixed. Input: scene_points: [numpy.ndarray, (N,3), numpy.float32] the scene points to detect voxel_size: [float] used for downsample Example usage: mfcdetector = ModelFreeCollisionDetector(scene_points, voxel_size=0.005) collision_mask = mfcdetector.detect(grasp_group, approach_dist=0.03) collision_mask, iou_list = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05, return_ious=True) collision_mask, empty_mask = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05, return_empty_grasp=True, empty_thresh=0.01) collision_mask, empty_mask, iou_list = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05, return_empty_grasp=True, empty_thresh=0.01, return_ious=True) """ def __init__(self, scene_points, voxel_size=0.005): self.finger_width = 0.01 self.finger_length = 0.06 self.voxel_size = voxel_size scene_cloud = o3d.geometry.PointCloud() scene_cloud.points = o3d.utility.Vector3dVector(scene_points) scene_cloud = scene_cloud.voxel_down_sample(voxel_size) self.scene_points = np.array(scene_cloud.points) def detect(self, grasp_group, approach_dist=0.03, collision_thresh=0.05, return_empty_grasp=False, empty_thresh=0.01, return_ious=False): """ Detect collision of grasps. Input: grasp_group: [GraspGroup, M grasps] the grasps to check approach_dist: [float] the distance for a gripper to move along approaching direction before grasping this shifting space requires no point either collision_thresh: [float] if global collision iou is greater than this threshold, a collision is detected return_empty_grasp: [bool] if True, return a mask to imply whether there are objects in a grasp empty_thresh: [float] if inner space iou is smaller than this threshold, a collision is detected only set when [return_empty_grasp] is True return_ious: [bool] if True, return global collision iou and part collision ious Output: collision_mask: [numpy.ndarray, (M,), numpy.bool] True implies collision [optional] empty_mask: [numpy.ndarray, (M,), numpy.bool] True implies empty grasp only returned when [return_empty_grasp] is True [optional] iou_list: list of [numpy.ndarray, (M,), numpy.float32] global and part collision ious, containing [global_iou, left_iou, right_iou, bottom_iou, shifting_iou] only returned when [return_ious] is True """ approach_dist = max(approach_dist, self.finger_width) T = grasp_group.translations R = grasp_group.rotation_matrices heights = grasp_group.heights[:,np.newaxis] depths = grasp_group.depths[:,np.newaxis] widths = grasp_group.widths[:,np.newaxis] targets = self.scene_points[np.newaxis,:,:] - T[:,np.newaxis,:] targets = np.matmul(targets, R) ## collision detection # height mask mask1 = ((targets[:,:,2] > -heights/2) & (targets[:,:,2] < heights/2)) # left finger mask mask2 = ((targets[:,:,0] > depths - self.finger_length) & (targets[:,:,0] < depths)) mask3 = (targets[:,:,1] > -(widths/2 + self.finger_width)) mask4 = (targets[:,:,1] < -widths/2) # right finger mask mask5 = (targets[:,:,1] < (widths/2 + self.finger_width)) mask6 = (targets[:,:,1] > widths/2) # bottom mask mask7 = ((targets[:,:,0] <= depths - self.finger_length)\ & (targets[:,:,0] > depths - self.finger_length - self.finger_width)) # shifting mask mask8 = ((targets[:,:,0] <= depths - self.finger_length - self.finger_width)\ & (targets[:,:,0] > depths - self.finger_length - self.finger_width - approach_dist)) # get collision mask of each point left_mask = (mask1 & mask2 & mask3 & mask4) right_mask = (mask1 & mask2 & mask5 & mask6) bottom_mask = (mask1 & mask3 & mask5 & mask7) shifting_mask = (mask1 & mask3 & mask5 & mask8) global_mask = (left_mask | right_mask | bottom_mask | shifting_mask) # calculate equivalant volume of each part left_right_volume = (heights * self.finger_length * self.finger_width / (self.voxel_size**3)).reshape(-1) bottom_volume = (heights * (widths+2*self.finger_width) * self.finger_width / (self.voxel_size**3)).reshape(-1) shifting_volume = (heights * (widths+2*self.finger_width) * approach_dist / (self.voxel_size**3)).reshape(-1) volume = left_right_volume*2 + bottom_volume + shifting_volume # get collision iou of each part global_iou = global_mask.sum(axis=1) / (volume+1e-6) # get collison mask collision_mask = (global_iou > collision_thresh) if not (return_empty_grasp or return_ious): return collision_mask ret_value = [collision_mask,] if return_empty_grasp: inner_mask = (mask1 & mask2 & (~mask4) & (~mask6)) inner_volume = (heights * self.finger_length * widths / (self.voxel_size**3)).reshape(-1) empty_mask = (inner_mask.sum(axis=-1)/inner_volume < empty_thresh) ret_value.append(empty_mask) if return_ious: left_iou = left_mask.sum(axis=1) / (left_right_volume+1e-6) right_iou = right_mask.sum(axis=1) / (left_right_volume+1e-6) bottom_iou = bottom_mask.sum(axis=1) / (bottom_volume+1e-6) shifting_iou = shifting_mask.sum(axis=1) / (shifting_volume+1e-6) ret_value.append([global_iou, left_iou, right_iou, bottom_iou, shifting_iou]) return ret_value ================================================ FILE: utils/data_utils.py ================================================ """ Tools for data processing. Author: chenxi-wang """ import numpy as np class CameraInfo(): """ Camera intrisics for point cloud creation. """ def __init__(self, width, height, fx, fy, cx, cy, scale): self.width = width self.height = height self.fx = fx self.fy = fy self.cx = cx self.cy = cy self.scale = scale def create_point_cloud_from_depth_image(depth, camera, organized=True): """ Generate point cloud using depth image only. Input: depth: [numpy.ndarray, (H,W), numpy.float32] depth image camera: [CameraInfo] camera intrinsics organized: bool whether to keep the cloud in image shape (H,W,3) Output: cloud: [numpy.ndarray, (H,W,3)/(H*W,3), numpy.float32] generated cloud, (H,W,3) for organized=True, (H*W,3) for organized=False """ assert (depth.shape[0] == camera.height and depth.shape[1] == camera.width) xmap = np.arange(camera.width) ymap = np.arange(camera.height) xmap, ymap = np.meshgrid(xmap, ymap) points_z = depth / camera.scale points_x = (xmap - camera.cx) * points_z / camera.fx points_y = (ymap - camera.cy) * points_z / camera.fy cloud = np.stack([points_x, points_y, points_z], axis=-1) if not organized: cloud = cloud.reshape([-1, 3]) return cloud def transform_point_cloud(cloud, transform, format='4x4'): """ Transform points to new coordinates with transformation matrix. Input: cloud: [np.ndarray, (N,3), np.float32] points in original coordinates transform: [np.ndarray, (3,3)/(3,4)/(4,4), np.float32] transformation matrix, could be rotation only or rotation+translation format: [string, '3x3'/'3x4'/'4x4'] the shape of transformation matrix '3x3' --> rotation matrix '3x4'/'4x4' --> rotation matrix + translation matrix Output: cloud_transformed: [np.ndarray, (N,3), np.float32] points in new coordinates """ if not (format == '3x3' or format == '4x4' or format == '3x4'): raise ValueError('Unknown transformation format, only support \'3x3\' or \'4x4\' or \'3x4\'.') if format == '3x3': cloud_transformed = np.dot(transform, cloud.T).T elif format == '4x4' or format == '3x4': ones = np.ones(cloud.shape[0])[:, np.newaxis] cloud_ = np.concatenate([cloud, ones], axis=1) cloud_transformed = np.dot(transform, cloud_.T).T cloud_transformed = cloud_transformed[:, :3] return cloud_transformed def compute_point_dists(A, B): """ Compute pair-wise point distances in two matrices. Input: A: [np.ndarray, (N,3), np.float32] point cloud A B: [np.ndarray, (M,3), np.float32] point cloud B Output: dists: [np.ndarray, (N,M), np.float32] distance matrix """ A = A[:, np.newaxis, :] B = B[np.newaxis, :, :] dists = np.linalg.norm(A - B, axis=-1) return dists def remove_invisible_grasp_points(cloud, grasp_points, pose, th=0.01): """ Remove invisible part of object model according to scene point cloud. Input: cloud: [np.ndarray, (N,3), np.float32] scene point cloud grasp_points: [np.ndarray, (M,3), np.float32] grasp point label in object coordinates pose: [np.ndarray, (4,4), np.float32] transformation matrix from object coordinates to world coordinates th: [float] if the minimum distance between a grasp point and the scene points is greater than outlier, the point will be removed Output: visible_mask: [np.ndarray, (M,), np.bool] mask to show the visible part of grasp points """ grasp_points_trans = transform_point_cloud(grasp_points, pose) dists = compute_point_dists(grasp_points_trans, cloud) min_dists = dists.min(axis=1) visible_mask = (min_dists < th) return visible_mask def get_workspace_mask(cloud, seg, trans=None, organized=True, outlier=0): """ Keep points in workspace as input. Input: cloud: [np.ndarray, (H,W,3), np.float32] scene point cloud seg: [np.ndarray, (H,W,), np.uint8] segmantation label of scene points trans: [np.ndarray, (4,4), np.float32] transformation matrix for scene points, default: None. organized: [bool] whether to keep the cloud in image shape (H,W,3) outlier: [float] if the distance between a point and workspace is greater than outlier, the point will be removed Output: workspace_mask: [np.ndarray, (H,W)/(H*W,), np.bool] mask to indicate whether scene points are in workspace """ if organized: h, w, _ = cloud.shape cloud = cloud.reshape([h * w, 3]) seg = seg.reshape(h * w) if trans is not None: cloud = transform_point_cloud(cloud, trans) foreground = cloud[seg > 0] xmin, ymin, zmin = foreground.min(axis=0) xmax, ymax, zmax = foreground.max(axis=0) mask_x = ((cloud[:, 0] > xmin - outlier) & (cloud[:, 0] < xmax + outlier)) mask_y = ((cloud[:, 1] > ymin - outlier) & (cloud[:, 1] < ymax + outlier)) mask_z = ((cloud[:, 2] > zmin - outlier) & (cloud[:, 2] < zmax + outlier)) workspace_mask = (mask_x & mask_y & mask_z) if organized: workspace_mask = workspace_mask.reshape([h, w]) return workspace_mask ================================================ FILE: utils/label_generation.py ================================================ """ Dynamically generate grasp labels during training. Author: chenxi-wang """ import os import sys import torch BASE_DIR = os.path.dirname(os.path.abspath(__file__)) ROOT_DIR = os.path.dirname(BASE_DIR) sys.path.append(ROOT_DIR) # sys.path.append(os.path.join(ROOT_DIR, 'knn')) from knn.knn_modules import knn from loss_utils import GRASP_MAX_WIDTH, batch_viewpoint_params_to_matrix, \ transform_point_cloud, generate_grasp_views def process_grasp_labels(end_points): """ Process labels according to scene points and object poses. """ seed_xyzs = end_points['xyz_graspable'] # (B, M_point, 3) batch_size, num_samples, _ = seed_xyzs.size() batch_grasp_points = [] batch_grasp_views_rot = [] batch_grasp_scores = [] batch_grasp_widths = [] for i in range(batch_size): seed_xyz = seed_xyzs[i] # (Ns, 3) poses = end_points['object_poses_list'][i] # [(3, 4),] # get merged grasp points for label computation grasp_points_merged = [] grasp_views_rot_merged = [] grasp_scores_merged = [] grasp_widths_merged = [] for obj_idx, pose in enumerate(poses): grasp_points = end_points['grasp_points_list'][i][obj_idx] # (Np, 3) grasp_scores = end_points['grasp_scores_list'][i][obj_idx] # (Np, V, A, D) grasp_widths = end_points['grasp_widths_list'][i][obj_idx] # (Np, V, A, D) _, V, A, D = grasp_scores.size() num_grasp_points = grasp_points.size(0) # generate and transform template grasp views grasp_views = generate_grasp_views(V).to(pose.device) # (V, 3) grasp_points_trans = transform_point_cloud(grasp_points, pose, '3x4') grasp_views_trans = transform_point_cloud(grasp_views, pose[:3, :3], '3x3') # generate and transform template grasp view rotation angles = torch.zeros(grasp_views.size(0), dtype=grasp_views.dtype, device=grasp_views.device) grasp_views_rot = batch_viewpoint_params_to_matrix(-grasp_views, angles) # (V, 3, 3) grasp_views_rot_trans = torch.matmul(pose[:3, :3], grasp_views_rot) # (V, 3, 3) # assign views grasp_views_ = grasp_views.transpose(0, 1).contiguous().unsqueeze(0) grasp_views_trans_ = grasp_views_trans.transpose(0, 1).contiguous().unsqueeze(0) view_inds = knn(grasp_views_trans_, grasp_views_, k=1).squeeze() - 1 grasp_views_rot_trans = torch.index_select(grasp_views_rot_trans, 0, view_inds) # (V, 3, 3) grasp_views_rot_trans = grasp_views_rot_trans.unsqueeze(0).expand(num_grasp_points, -1, -1, -1) # (Np, V, 3, 3) grasp_scores = torch.index_select(grasp_scores, 1, view_inds) # (Np, V, A, D) grasp_widths = torch.index_select(grasp_widths, 1, view_inds) # (Np, V, A, D) # add to list grasp_points_merged.append(grasp_points_trans) grasp_views_rot_merged.append(grasp_views_rot_trans) grasp_scores_merged.append(grasp_scores) grasp_widths_merged.append(grasp_widths) grasp_points_merged = torch.cat(grasp_points_merged, dim=0) # (Np', 3) grasp_views_rot_merged = torch.cat(grasp_views_rot_merged, dim=0) # (Np', V, 3, 3) grasp_scores_merged = torch.cat(grasp_scores_merged, dim=0) # (Np', V, A, D) grasp_widths_merged = torch.cat(grasp_widths_merged, dim=0) # (Np', V, A, D) # compute nearest neighbors seed_xyz_ = seed_xyz.transpose(0, 1).contiguous().unsqueeze(0) # (1, 3, Ns) grasp_points_merged_ = grasp_points_merged.transpose(0, 1).contiguous().unsqueeze(0) # (1, 3, Np') nn_inds = knn(grasp_points_merged_, seed_xyz_, k=1).squeeze() - 1 # (Ns) # assign anchor points to real points grasp_points_merged = torch.index_select(grasp_points_merged, 0, nn_inds) # (Ns, 3) grasp_views_rot_merged = torch.index_select(grasp_views_rot_merged, 0, nn_inds) # (Ns, V, 3, 3) grasp_scores_merged = torch.index_select(grasp_scores_merged, 0, nn_inds) # (Ns, V, A, D) grasp_widths_merged = torch.index_select(grasp_widths_merged, 0, nn_inds) # (Ns, V, A, D) # add to batch batch_grasp_points.append(grasp_points_merged) batch_grasp_views_rot.append(grasp_views_rot_merged) batch_grasp_scores.append(grasp_scores_merged) batch_grasp_widths.append(grasp_widths_merged) batch_grasp_points = torch.stack(batch_grasp_points, 0) # (B, Ns, 3) batch_grasp_views_rot = torch.stack(batch_grasp_views_rot, 0) # (B, Ns, V, 3, 3) batch_grasp_scores = torch.stack(batch_grasp_scores, 0) # (B, Ns, V, A, D) batch_grasp_widths = torch.stack(batch_grasp_widths, 0) # (B, Ns, V, A, D) # compute view graspness view_u_threshold = 0.6 view_grasp_num = 48 batch_grasp_view_valid_mask = (batch_grasp_scores <= view_u_threshold) & (batch_grasp_scores > 0) # (B, Ns, V, A, D) batch_grasp_view_valid = batch_grasp_view_valid_mask.float() batch_grasp_view_graspness = torch.sum(torch.sum(batch_grasp_view_valid, dim=-1), dim=-1) / view_grasp_num # (B, Ns, V) view_graspness_min, _ = torch.min(batch_grasp_view_graspness, dim=-1) # (B, Ns) view_graspness_max, _ = torch.max(batch_grasp_view_graspness, dim=-1) view_graspness_max = view_graspness_max.unsqueeze(-1).expand(-1, -1, 300) # (B, Ns, V) view_graspness_min = view_graspness_min.unsqueeze(-1).expand(-1, -1, 300) # same shape as batch_grasp_view_graspness batch_grasp_view_graspness = (batch_grasp_view_graspness - view_graspness_min) / (view_graspness_max - view_graspness_min + 1e-5) # process scores label_mask = (batch_grasp_scores > 0) & (batch_grasp_widths <= GRASP_MAX_WIDTH) # (B, Ns, V, A, D) batch_grasp_scores[~label_mask] = 0 end_points['batch_grasp_point'] = batch_grasp_points end_points['batch_grasp_view_rot'] = batch_grasp_views_rot end_points['batch_grasp_score'] = batch_grasp_scores end_points['batch_grasp_width'] = batch_grasp_widths end_points['batch_grasp_view_graspness'] = batch_grasp_view_graspness return end_points def match_grasp_view_and_label(end_points): """ Slice grasp labels according to predicted views. """ top_view_inds = end_points['grasp_top_view_inds'] # (B, Ns) template_views_rot = end_points['batch_grasp_view_rot'] # (B, Ns, V, 3, 3) grasp_scores = end_points['batch_grasp_score'] # (B, Ns, V, A, D) grasp_widths = end_points['batch_grasp_width'] # (B, Ns, V, A, D, 3) B, Ns, V, A, D = grasp_scores.size() top_view_inds_ = top_view_inds.view(B, Ns, 1, 1, 1).expand(-1, -1, -1, 3, 3) top_template_views_rot = torch.gather(template_views_rot, 2, top_view_inds_).squeeze(2) top_view_inds_ = top_view_inds.view(B, Ns, 1, 1, 1).expand(-1, -1, -1, A, D) top_view_grasp_scores = torch.gather(grasp_scores, 2, top_view_inds_).squeeze(2) top_view_grasp_widths = torch.gather(grasp_widths, 2, top_view_inds_).squeeze(2) u_max = top_view_grasp_scores.max() po_mask = top_view_grasp_scores > 0 po_mask_num = torch.sum(po_mask) if po_mask_num > 0: u_min = top_view_grasp_scores[po_mask].min() top_view_grasp_scores[po_mask] = torch.log(u_max / top_view_grasp_scores[po_mask]) / (torch.log(u_max / u_min) + 1e-6) end_points['batch_grasp_score'] = top_view_grasp_scores # (B, Ns, A, D) end_points['batch_grasp_width'] = top_view_grasp_widths # (B, Ns, A, D) return top_template_views_rot, end_points ================================================ FILE: utils/loss_utils.py ================================================ """ Tools for loss computation. Author: chenxi-wang """ import torch import numpy as np GRASP_MAX_WIDTH = 0.1 GRASPNESS_THRESHOLD = 0.1 NUM_VIEW = 300 NUM_ANGLE = 12 NUM_DEPTH = 4 M_POINT = 1024 def transform_point_cloud(cloud, transform, format='4x4'): """ Transform points to new coordinates with transformation matrix. Input: cloud: [torch.FloatTensor, (N,3)] points in original coordinates transform: [torch.FloatTensor, (3,3)/(3,4)/(4,4)] transformation matrix, could be rotation only or rotation+translation format: [string, '3x3'/'3x4'/'4x4'] the shape of transformation matrix '3x3' --> rotation matrix '3x4'/'4x4' --> rotation matrix + translation matrix Output: cloud_transformed: [torch.FloatTensor, (N,3)] points in new coordinates """ if not (format == '3x3' or format == '4x4' or format == '3x4'): raise ValueError('Unknown transformation format, only support \'3x3\' or \'4x4\' or \'3x4\'.') if format == '3x3': cloud_transformed = torch.matmul(transform, cloud.T).T elif format == '4x4' or format == '3x4': ones = cloud.new_ones(cloud.size(0), device=cloud.device).unsqueeze(-1) cloud_ = torch.cat([cloud, ones], dim=1) cloud_transformed = torch.matmul(transform, cloud_.T).T cloud_transformed = cloud_transformed[:, :3] return cloud_transformed def generate_grasp_views(N=300, phi=(np.sqrt(5) - 1) / 2, center=np.zeros(3), r=1): """ View sampling on a unit sphere using Fibonacci lattices. Ref: https://arxiv.org/abs/0912.4540 Input: N: [int] number of sampled views phi: [float] constant for view coordinate calculation, different phi's bring different distributions, default: (sqrt(5)-1)/2 center: [np.ndarray, (3,), np.float32] sphere center r: [float] sphere radius Output: views: [torch.FloatTensor, (N,3)] sampled view coordinates """ views = [] for i in range(N): zi = (2 * i + 1) / N - 1 xi = np.sqrt(1 - zi ** 2) * np.cos(2 * i * np.pi * phi) yi = np.sqrt(1 - zi ** 2) * np.sin(2 * i * np.pi * phi) views.append([xi, yi, zi]) views = r * np.array(views) + center return torch.from_numpy(views.astype(np.float32)) def batch_viewpoint_params_to_matrix(batch_towards, batch_angle): """ Transform approach vectors and in-plane rotation angles to rotation matrices. Input: batch_towards: [torch.FloatTensor, (N,3)] approach vectors in batch batch_angle: [torch.floatTensor, (N,)] in-plane rotation angles in batch Output: batch_matrix: [torch.floatTensor, (N,3,3)] rotation matrices in batch """ axis_x = batch_towards ones = torch.ones(axis_x.shape[0], dtype=axis_x.dtype, device=axis_x.device) zeros = torch.zeros(axis_x.shape[0], dtype=axis_x.dtype, device=axis_x.device) axis_y = torch.stack([-axis_x[:, 1], axis_x[:, 0], zeros], dim=-1) mask_y = (torch.norm(axis_y, dim=-1) == 0) axis_y[mask_y, 1] = 1 axis_x = axis_x / torch.norm(axis_x, dim=-1, keepdim=True) axis_y = axis_y / torch.norm(axis_y, dim=-1, keepdim=True) axis_z = torch.cross(axis_x, axis_y) sin = torch.sin(batch_angle) cos = torch.cos(batch_angle) R1 = torch.stack([ones, zeros, zeros, zeros, cos, -sin, zeros, sin, cos], dim=-1) R1 = R1.reshape([-1, 3, 3]) R2 = torch.stack([axis_x, axis_y, axis_z], dim=-1) batch_matrix = torch.matmul(R2, R1) return batch_matrix def huber_loss(error, delta=1.0): """ Args: error: Torch tensor (d1,d2,...,dk) Returns: loss: Torch tensor (d1,d2,...,dk) x = error = pred - gt or dist(pred,gt) 0.5 * |x|^2 if |x|<=d 0.5 * d^2 + d * (|x|-d) if |x|>d Author: Charles R. Qi Ref: https://github.com/charlesq34/frustum-pointnets/blob/master/models/model_util.py """ abs_error = torch.abs(error) quadratic = torch.clamp(abs_error, max=delta) linear = (abs_error - quadratic) loss = 0.5 * quadratic ** 2 + delta * linear return loss