Showing preview only (213K chars total). Download the full file or copy to clipboard to get everything.
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__ | AP<sub>0.8</sub> | AP<sub>0.4</sub> | __AP__ | AP<sub>0.8</sub> | AP<sub>0.4</sub> | __AP__ | AP<sub>0.8</sub> | AP<sub>0.4</sub> |
| 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<query_width;query_idx++)
{
for(int ref_idx = 0;ref_idx < ref_width;ref_idx++)
{
dist_dev[query_idx * ref_width + ref_idx] = 0;
for(int hi=0;hi<height;hi++)
dist_dev[query_idx * ref_width + ref_idx] += (ref_dev[hi * ref_width + ref_idx] - query_dev[hi * query_width + query_idx]) * (ref_dev[hi * ref_width + ref_idx] - query_dev[hi * query_width + query_idx]);
}
}
float temp_value;
long temp_idx;
// sort the distance and get the index
for(int query_idx = 0;query_idx<query_width;query_idx++)
{
for(int i = 0;i < ref_width;i++)
{
ind_buf[i] = i+1;
}
for(int i = 0;i < ref_width;i++)
for(int j = 0;j < ref_width - i - 1;j++)
{
if(dist_dev[query_idx * ref_width + j] > 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 <torch/extension.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);
================================================
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 <chrischoy@ai.stanford.edu> 12/23/2016
*/
// Includes
#include <cstdio>
#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 (xIndex<width){
// Pointer shift, initialization, and max value
p_dist = dist + xIndex;
p_ind = ind + xIndex;
max_dist = p_dist[0];
p_ind[0] = 1;
// Part 1 : sort kth firt elementZ
for (l=1; l<k; l++){
curr_row = l * width;
curr_dist = p_dist[curr_row];
if (curr_dist<max_dist){
i=l-1;
for (int a=0; a<l-1; a++){
if (p_dist[a*width]>curr_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; l<height; l++){
curr_dist = p_dist[l*width];
if (curr_dist<max_dist){
i=k-1;
for (int a=0; a<k-1; a++){
if (p_dist[a*width]>curr_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<width && yIndex<k)
dist[yIndex*width + xIndex] = sqrt(dist[yIndex*width + xIndex]);
}
//-----------------------------------------------------------------------------------------------//
// K-th NEAREST NEIGHBORS //
//-----------------------------------------------------------------------------------------------//
/**
* K nearest neighbor algorithm
* - Initialize CUDA
* - Allocate device memory
* - Copy point sets (reference and query points) from host to device memory
* - Compute the distances + indexes to the k nearest neighbors for each query point
* - Copy distances from device to host memory
*
* @param ref_host reference points ; pointer to linear matrix
* @param ref_nb number of reference points ; width of the matrix
* @param query_host query points ; pointer to linear matrix
* @param query_nb number of query points ; width of the matrix
* @param dim dimension of points ; height of the matrices
* @param k number of neighbor to consider
* @param dist_host distances to k nearest neighbors ; pointer to linear matrix
* @param dist_host indexes of the k nearest neighbors ; pointer to linear matrix
*
*/
void knn_device(float* ref_dev, int ref_nb, float* query_dev, int query_nb,
int dim, int k, float* dist_dev, long* ind_dev, cudaStream_t stream){
// Grids and threads
dim3 g_16x16(query_nb/16, ref_nb/16, 1);
dim3 t_16x16(16, 16, 1);
if (query_nb%16 != 0) g_16x16.x += 1;
if (ref_nb %16 != 0) g_16x16.y += 1;
//
dim3 g_256x1(query_nb/256, 1, 1);
dim3 t_256x1(256, 1, 1);
if (query_nb%256 != 0) g_256x1.x += 1;
dim3 g_k_16x16(query_nb/16, k/16, 1);
dim3 t_k_16x16(16, 16, 1);
if (query_nb%16 != 0) g_k_16x16.x += 1;
if (k %16 != 0) g_k_16x16.y += 1;
// Kernel 1: Compute all the distances
cuComputeDistanceGlobal<<<g_16x16, t_16x16, 0, stream>>>(ref_dev, ref_nb, query_dev, query_nb, dim, dist_dev);
// Kernel 2: Sort each column
cuInsertionSort<<<g_256x1, t_256x1, 0, stream>>>(dist_dev, ind_dev, query_nb, ref_nb, k);
// Kernel 3: Compute square root of k first elements
// cuParallelSqrt<<<g_k_16x16,t_k_16x16, 0, stream>>>(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 <torch/extension.h>
#include <THC/THC.h>
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 <THC/THC.h>
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>();
float *query_dev = query.data<float>();
long *idx_dev = idx.data<long>();
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 <torch/extension.h>
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 <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <cmath>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#define TOTAL_THREADS 512
inline int opt_n_threads(int work_size) {
const int pow_2 = std::log(static_cast<double>(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 <torch/extension.h>
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 <torch/extension.h>
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 <torch/extension.h>
#include <vector>
std::vector<at::Tensor> 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 <torch/extension.h>
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 <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>
#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<float>(),
xyz.data<float>(), idx.data<int>());
} 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 <math.h>
#include <stdio.h>
#include <stdlib.h>
#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, opt_n_threads(m), 0, stream>>>(
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<float>(),
xyz.data<float>(), rot.data<float>(), idx.data<int>());
} else {
TORCH_CHECK(false, "CPU not supported");
}
return idx;
}
================================================
FILE: pointnet2/_ext_src/src/cylinder_query_gpu.cu
================================================
// Author: chenxi-wang
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#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, opt_n_threads(m), 0, stream>>>(
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<float>(),
idx.data<int>(), output.data<float>());
} 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<float>(), idx.data<int>(), output.data<float>());
} 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 <stdio.h>
#include <stdlib.h>
#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, opt_block_config(npoints, c), 0, stream>>>(
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, opt_block_config(npoints, c), 0, stream>>>(
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<at::Tensor> 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<float>(), knows.data<float>(),
dist2.data<float>(), idx.data<int>());
} 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<float>(), idx.data<int>(), weight.data<float>(),
output.data<float>());
} 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<float>(), idx.data<int>(), weight.data<float>(),
output.data<float>());
} 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 <math.h>
#include <stdio.h>
#include <stdlib.h>
#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, opt_n_threads(n), 0, stream>>>(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, opt_block_config(n, c), 0, stream>>>(
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, opt_block_config(n, c), 0, stream>>>(
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<float>(),
idx.data<int>(), output.data<float>());
} 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<float>(),
idx.data<int>(), output.data<float>());
} 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<float>(),
tmp.data<float>(), output.data<int>());
} 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 <stdio.h>
#include <stdlib.h>
#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<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
at::cuda::getCurrentCUDAStream()>>>(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<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
at::cuda::getCurrentCUDAStream()>>>(
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 <unsigned int block_size>
__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_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 256:
furthest_point_sampling_kernel<256>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 128:
furthest_point_sampling_kernel<128>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 64:
furthest_point_sampling_kernel<64>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 32:
furthest_point_sampling_kernel<32>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 16:
furthest_point_sampling_kernel<16>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 8:
furthest_point_sampling_kernel<8>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 4:
furthest_point_sampling_kernel<4>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 2:
furthest_point_sampling_kernel<2>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 1:
furthest_point_sampling_kernel<1>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
default:
furthest_point_sampling_kernel<512>
<<<b, n_threads, 0, stream>>>(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
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
SYMBOL INDEX (195 symbols across 29 files)
FILE: dataset/graspnet_dataset.py
class GraspNetDataset (line 18) | class GraspNetDataset(Dataset):
method __init__ (line 19) | def __init__(self, root, grasp_labels=None, camera='kinect', split='tr...
method scene_list (line 65) | def scene_list(self):
method __len__ (line 68) | def __len__(self):
method augment_data (line 71) | def augment_data(self, point_clouds, object_poses_list):
method __getitem__ (line 93) | def __getitem__(self, index):
method get_data (line 99) | def get_data(self, index, return_raw_cloud=False):
method get_data_label (line 145) | def get_data_label(self, index):
function load_grasp_labels (line 226) | def load_grasp_labels(root):
function minkowski_collate_fn (line 237) | def minkowski_collate_fn(list_data):
FILE: dataset/simplify_dataset.py
function simplify_grasp_labels (line 10) | def simplify_grasp_labels(root, save_path):
FILE: infer_vis_grasp.py
function data_process (line 43) | def data_process():
function my_worker_init_fn (line 87) | def my_worker_init_fn(worker_id):
function inference (line 92) | def inference(data_input):
FILE: knn/knn_modules.py
function knn (line 9) | def knn(ref, query, k=1):
FILE: knn/setup.py
function get_extensions (line 16) | def get_extensions():
FILE: knn/src/cpu/knn_cpu.cpp
function knn_cpu (line 4) | void knn_cpu(float* ref_dev, int ref_width, float* query_dev, int query_...
FILE: knn/src/knn.h
function knn (line 12) | int knn(at::Tensor& ref, at::Tensor& query, at::Tensor& idx)
FILE: knn/src/vision.cpp
function PYBIND11_MODULE (line 3) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: models/backbone_resunet14.py
class MinkUNetBase (line 6) | class MinkUNetBase(ResNetBase):
method __init__ (line 18) | def __init__(self, in_channels, out_channels, D=3):
method network_initialization (line 21) | def network_initialization(self, in_channels, out_channels, D):
method forward (line 93) | def forward(self, x):
class MinkUNet14 (line 154) | class MinkUNet14(MinkUNetBase):
class MinkUNet18 (line 159) | class MinkUNet18(MinkUNetBase):
class MinkUNet34 (line 164) | class MinkUNet34(MinkUNetBase):
class MinkUNet50 (line 169) | class MinkUNet50(MinkUNetBase):
class MinkUNet101 (line 174) | class MinkUNet101(MinkUNetBase):
class MinkUNet14A (line 179) | class MinkUNet14A(MinkUNet14):
class MinkUNet14B (line 183) | class MinkUNet14B(MinkUNet14):
class MinkUNet14C (line 187) | class MinkUNet14C(MinkUNet14):
class MinkUNet14Dori (line 191) | class MinkUNet14Dori(MinkUNet14):
class MinkUNet14E (line 195) | class MinkUNet14E(MinkUNet14):
class MinkUNet14D (line 199) | class MinkUNet14D(MinkUNet14):
class MinkUNet18A (line 203) | class MinkUNet18A(MinkUNet18):
class MinkUNet18B (line 207) | class MinkUNet18B(MinkUNet18):
class MinkUNet18D (line 211) | class MinkUNet18D(MinkUNet18):
class MinkUNet34A (line 215) | class MinkUNet34A(MinkUNet34):
class MinkUNet34B (line 219) | class MinkUNet34B(MinkUNet34):
class MinkUNet34C (line 223) | class MinkUNet34C(MinkUNet34):
FILE: models/graspnet.py
class GraspNet (line 23) | class GraspNet(nn.Module):
method __init__ (line 24) | def __init__(self, cylinder_radius=0.05, seed_feat_dim=512, is_trainin...
method forward (line 39) | def forward(self, end_points):
function pred_decode (line 96) | def pred_decode(end_points):
FILE: models/loss.py
function get_loss (line 5) | def get_loss(end_points):
function compute_objectness_loss (line 16) | def compute_objectness_loss(end_points):
function compute_graspness_loss (line 32) | def compute_graspness_loss(end_points):
function compute_view_graspness_loss (line 52) | def compute_view_graspness_loss(end_points):
function compute_score_loss (line 61) | def compute_score_loss(end_points):
function compute_width_loss (line 71) | def compute_width_loss(end_points):
FILE: models/modules.py
class GraspableNet (line 16) | class GraspableNet(nn.Module):
method __init__ (line 17) | def __init__(self, seed_feature_dim):
method forward (line 22) | def forward(self, seed_features, end_points):
class ApproachNet (line 29) | class ApproachNet(nn.Module):
method __init__ (line 30) | def __init__(self, num_view, seed_feature_dim, is_training=True):
method forward (line 38) | def forward(self, seed_features, end_points):
class CloudCrop (line 76) | class CloudCrop(nn.Module):
method __init__ (line 77) | def __init__(self, nsample, seed_feature_dim, cylinder_radius=0.05, hm...
method forward (line 88) | def forward(self, seed_xyz_graspable, seed_features_graspable, vp_rot):
class SWADNet (line 97) | class SWADNet(nn.Module):
method __init__ (line 98) | def __init__(self, num_angle, num_depth):
method forward (line 106) | def forward(self, vp_features, end_points):
FILE: models/resnet.py
class ResNetBase (line 12) | class ResNetBase(nn.Module):
method __init__ (line 18) | def __init__(self, in_channels, out_channels, D=3):
method network_initialization (line 26) | def network_initialization(self, in_channels, out_channels, D):
method weight_initialization (line 64) | def weight_initialization(self):
method _make_layer (line 73) | def _make_layer(self, block, planes, blocks, stride=1, dilation=1, bn_...
method forward (line 107) | def forward(self, x: ME.SparseTensor):
class ResNet14 (line 118) | class ResNet14(ResNetBase):
class ResNet18 (line 123) | class ResNet18(ResNetBase):
class ResNet34 (line 128) | class ResNet34(ResNetBase):
class ResNet50 (line 133) | class ResNet50(ResNetBase):
class ResNet101 (line 138) | class ResNet101(ResNetBase):
class ResFieldNetBase (line 143) | class ResFieldNetBase(ResNetBase):
method network_initialization (line 144) | def network_initialization(self, in_channels, out_channels, D):
method forward (line 168) | def forward(self, x: ME.TensorField):
class ResFieldNet14 (line 174) | class ResFieldNet14(ResFieldNetBase):
class ResFieldNet18 (line 179) | class ResFieldNet18(ResFieldNetBase):
class ResFieldNet34 (line 184) | class ResFieldNet34(ResFieldNetBase):
class ResFieldNet50 (line 189) | class ResFieldNet50(ResFieldNetBase):
class ResFieldNet101 (line 194) | class ResFieldNet101(ResFieldNetBase):
FILE: pointnet2/_ext_src/include/cuda_utils.h
function opt_n_threads (line 20) | inline int opt_n_threads(int work_size) {
function dim3 (line 26) | inline dim3 opt_block_config(int x, int y) {
FILE: pointnet2/_ext_src/src/ball_query.cpp
function ball_query (line 13) | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float ra...
FILE: pointnet2/_ext_src/src/bindings.cpp
function PYBIND11_MODULE (line 12) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: pointnet2/_ext_src/src/cylinder_query.cpp
function cylinder_query (line 10) | at::Tensor cylinder_query(at::Tensor new_xyz, at::Tensor xyz, at::Tensor...
FILE: pointnet2/_ext_src/src/group_points.cpp
function group_points (line 17) | at::Tensor group_points(at::Tensor points, at::Tensor idx) {
function group_points_grad (line 42) | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const ...
FILE: pointnet2/_ext_src/src/interpolate.cpp
function three_nn (line 19) | std::vector<at::Tensor> three_nn(at::Tensor unknowns, at::Tensor knows) {
function three_interpolate (line 47) | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx,
function three_interpolate_grad (line 76) | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx,
FILE: pointnet2/_ext_src/src/sampling.cpp
function gather_points (line 20) | at::Tensor gather_points(at::Tensor points, at::Tensor idx) {
function gather_points_grad (line 45) | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx,
function furthest_point_sampling (line 70) | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) {
FILE: pointnet2/pointnet2_modules.py
class _PointnetSAModuleBase (line 26) | class _PointnetSAModuleBase(nn.Module):
method __init__ (line 28) | def __init__(self):
method forward (line 34) | def forward(self, xyz: torch.Tensor,
class PointnetSAModuleMSG (line 78) | class PointnetSAModuleMSG(_PointnetSAModuleBase):
method __init__ (line 95) | def __init__(
class PointnetSAModule (line 127) | class PointnetSAModule(PointnetSAModuleMSG):
method __init__ (line 144) | def __init__(
class PointnetSAModuleVotes (line 164) | class PointnetSAModuleVotes(nn.Module):
method __init__ (line 168) | def __init__(
method forward (line 210) | def forward(self, xyz: torch.Tensor,
class PointnetSAModuleMSGVotes (line 274) | class PointnetSAModuleMSGVotes(nn.Module):
method __init__ (line 278) | def __init__(
method forward (line 309) | def forward(self, xyz: torch.Tensor,
class PointnetFPModule (line 356) | class PointnetFPModule(nn.Module):
method __init__ (line 367) | def __init__(self, *, mlp: List[int], bn: bool = True):
method forward (line 371) | def forward(
class PointnetLFPModuleMSG (line 418) | class PointnetLFPModuleMSG(nn.Module):
method __init__ (line 422) | def __init__(
method forward (line 454) | def forward(self, xyz2: torch.Tensor, xyz1: torch.Tensor,
FILE: pointnet2/pointnet2_utils.py
class RandomDropout (line 40) | class RandomDropout(nn.Module):
method __init__ (line 41) | def __init__(self, p=0.5, inplace=False):
method forward (line 46) | def forward(self, X):
class FurthestPointSampling (line 51) | class FurthestPointSampling(Function):
method forward (line 53) | def forward(ctx, xyz, npoint):
method backward (line 74) | def backward(xyz, a=None):
class GatherOperation (line 81) | class GatherOperation(Function):
method forward (line 83) | def forward(ctx, features, idx):
method backward (line 108) | def backward(ctx, grad_out):
class ThreeNN (line 118) | class ThreeNN(Function):
method forward (line 120) | def forward(ctx, unknown, known):
method backward (line 143) | def backward(ctx, a=None, b=None):
class ThreeInterpolate (line 150) | class ThreeInterpolate(Function):
method forward (line 152) | def forward(ctx, features, idx, weight):
method backward (line 178) | def backward(ctx, grad_out):
class GroupingOperation (line 207) | class GroupingOperation(Function):
method forward (line 209) | def forward(ctx, features, idx):
method backward (line 233) | def backward(ctx, grad_out):
class BallQuery (line 258) | class BallQuery(Function):
method forward (line 260) | def forward(ctx, radius, nsample, xyz, new_xyz):
method backward (line 283) | def backward(ctx, a=None):
class QueryAndGroup (line 290) | class QueryAndGroup(nn.Module):
method __init__ (line 302) | def __init__(self, radius, nsample, use_xyz=True, ret_grouped_xyz=Fals...
method forward (line 313) | def forward(self, xyz, new_xyz, features=None):
class GroupAll (line 375) | class GroupAll(nn.Module):
method __init__ (line 383) | def __init__(self, use_xyz=True, ret_grouped_xyz=False):
method forward (line 388) | def forward(self, xyz, new_xyz, features=None):
class CylinderQuery (line 424) | class CylinderQuery(Function):
method forward (line 426) | def forward(ctx, radius, hmin, hmax, nsample, xyz, new_xyz, rot):
method backward (line 454) | def backward(ctx, a=None):
class CylinderQueryAndGroup (line 461) | class CylinderQueryAndGroup(nn.Module):
method __init__ (line 475) | def __init__(self, radius, hmin, hmax, nsample, use_xyz=True, ret_grou...
method forward (line 487) | def forward(self, xyz, new_xyz, rot, features=None):
FILE: pointnet2/pytorch_utils.py
class SharedMLP (line 11) | class SharedMLP(nn.Sequential):
method __init__ (line 13) | def __init__(
class _BNBase (line 39) | class _BNBase(nn.Sequential):
method __init__ (line 41) | def __init__(self, in_size, batch_norm=None, name=""):
class BatchNorm1d (line 49) | class BatchNorm1d(_BNBase):
method __init__ (line 51) | def __init__(self, in_size: int, *, name: str = ""):
class BatchNorm2d (line 55) | class BatchNorm2d(_BNBase):
method __init__ (line 57) | def __init__(self, in_size: int, name: str = ""):
class BatchNorm3d (line 61) | class BatchNorm3d(_BNBase):
method __init__ (line 63) | def __init__(self, in_size: int, name: str = ""):
class _ConvBase (line 67) | class _ConvBase(nn.Sequential):
method __init__ (line 69) | def __init__(
class Conv1d (line 123) | class Conv1d(_ConvBase):
method __init__ (line 125) | def __init__(
class Conv2d (line 157) | class Conv2d(_ConvBase):
method __init__ (line 159) | def __init__(
class Conv3d (line 191) | class Conv3d(_ConvBase):
method __init__ (line 193) | def __init__(
class FC (line 225) | class FC(nn.Sequential):
method __init__ (line 227) | def __init__(
function set_bn_momentum_default (line 262) | def set_bn_momentum_default(bn_momentum):
class BNMomentumScheduler (line 271) | class BNMomentumScheduler(object):
method __init__ (line 273) | def __init__(
method step (line 291) | def step(self, epoch=None):
FILE: test.py
function my_worker_init_fn (line 42) | def my_worker_init_fn(worker_id):
function inference (line 47) | def inference():
function evaluate (line 108) | def evaluate(dump_dir):
FILE: train.py
function log_string (line 46) | def log_string(out_str):
function my_worker_init_fn (line 53) | def my_worker_init_fn(worker_id):
function get_current_lr (line 83) | def get_current_lr(epoch):
function adjust_learning_rate (line 89) | def adjust_learning_rate(optimizer, epoch):
function train_one_epoch (line 95) | def train_one_epoch():
function train (line 130) | def train(start_epoch):
FILE: utils/collision_detector.py
class ModelFreeCollisionDetector (line 10) | class ModelFreeCollisionDetector():
method __init__ (line 28) | def __init__(self, scene_points, voxel_size=0.005):
method detect (line 37) | def detect(self, grasp_group, approach_dist=0.03, collision_thresh=0.0...
FILE: utils/data_utils.py
class CameraInfo (line 8) | class CameraInfo():
method __init__ (line 11) | def __init__(self, width, height, fx, fy, cx, cy, scale):
function create_point_cloud_from_depth_image (line 21) | def create_point_cloud_from_depth_image(depth, camera, organized=True):
function transform_point_cloud (line 49) | def transform_point_cloud(cloud, transform, format='4x4'):
function compute_point_dists (line 78) | def compute_point_dists(A, B):
function remove_invisible_grasp_points (line 97) | def remove_invisible_grasp_points(cloud, grasp_points, pose, th=0.01):
function get_workspace_mask (line 121) | def get_workspace_mask(cloud, seg, trans=None, organized=True, outlier=0):
FILE: utils/label_generation.py
function process_grasp_labels (line 19) | def process_grasp_labels(end_points):
function match_grasp_view_and_label (line 119) | def match_grasp_view_and_label(end_points):
FILE: utils/loss_utils.py
function transform_point_cloud (line 16) | def transform_point_cloud(cloud, transform, format='4x4'):
function generate_grasp_views (line 45) | def generate_grasp_views(N=300, phi=(np.sqrt(5) - 1) / 2, center=np.zero...
function batch_viewpoint_params_to_matrix (line 73) | def batch_viewpoint_params_to_matrix(batch_towards, batch_angle):
function huber_loss (line 104) | def huber_loss(error, delta=1.0):
Condensed preview — 53 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (212K chars).
[
{
"path": ".gitignore",
"chars": 195,
"preview": "**/__pycache__/**\n*.ipynb\n**/.ipynb_checkpoints/**\n*.npy\n*.npz\n**/.vscode/**\n**/grasp_label*/**\n**/log*/**\n**/dump*/**\n*"
},
{
"path": "LICENSE",
"chars": 11784,
"preview": "GRASPNET-BASELINE\nSOFTWARE LICENSE AGREEMENT\nACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY\n\nBY USI"
},
{
"path": "README.md",
"chars": 3767,
"preview": "# GraspNet graspness\nMy implementation of paper \"Graspness Discovery in Clutters for Fast and Accurate Grasp Detection\" "
},
{
"path": "command_test.sh",
"chars": 226,
"preview": "CUDA_VISIBLE_DEVICES=4 python test.py --camera kinect --dump_dir logs/log_kn/dump_epoch10 --checkpoint_path logs/log_kn/"
},
{
"path": "command_train.sh",
"chars": 170,
"preview": "CUDA_VISIBLE_DEVICES=4 python train.py --camera kinect --log_dir logs/log_kn --batch_size 4 --learning_rate 0.001 --mode"
},
{
"path": "dataset/generate_graspness.py",
"chars": 6795,
"preview": "import numpy as np\nimport os\nfrom PIL import Image\nimport scipy.io as scio\nimport sys\nROOT_DIR = os.path.dirname(os.path"
},
{
"path": "dataset/graspnet_dataset.py",
"chars": 11541,
"preview": "\"\"\" GraspNet dataset processing.\n Author: chenxi-wang\n\"\"\"\n\nimport os\nimport numpy as np\nimport scipy.io as scio\nfrom "
},
{
"path": "dataset/simplify_dataset.py",
"chars": 1657,
"preview": "import numpy as np\nimport os\nimport argparse\n\n\nparser = argparse.ArgumentParser()\nparser.add_argument('--dataset_root', "
},
{
"path": "dataset/vis_graspness.py",
"chars": 2138,
"preview": "import open3d as o3d\nimport scipy.io as scio\nfrom PIL import Image\nimport os\nimport numpy as np\nimport sys\nROOT_DIR = os"
},
{
"path": "infer_vis_grasp.py",
"chars": 6503,
"preview": "import os\nimport sys\nimport numpy as np\nimport argparse\nfrom PIL import Image\nimport time\nimport scipy.io as scio\nimport"
},
{
"path": "knn/knn_modules.py",
"chars": 481,
"preview": "import unittest\nimport gc\nimport operator as op\nimport functools\nimport torch\nfrom torch.autograd import Variable, Funct"
},
{
"path": "knn/setup.py",
"chars": 1871,
"preview": "#!/usr/bin/env python\n\nimport glob\nimport os\n\nimport torch\nfrom setuptools import find_packages\nfrom setuptools import s"
},
{
"path": "knn/src/cpu/knn_cpu.cpp",
"chars": 1810,
"preview": "#include \"cpu/vision.h\"\n\n\nvoid knn_cpu(float* ref_dev, int ref_width, float* query_dev, int query_width,\n int height,"
},
{
"path": "knn/src/cpu/vision.h",
"chars": 196,
"preview": "#pragma once\n#include <torch/extension.h>\n\nvoid knn_cpu(float* ref_dev, int ref_width,\n float* query_dev, int query_w"
},
{
"path": "knn/src/cuda/knn.cu",
"chars": 8739,
"preview": "/** Modifed version of knn-CUDA from https://github.com/vincentfpgarcia/kNN-CUDA\n * The modifications are\n * remove"
},
{
"path": "knn/src/cuda/vision.h",
"chars": 226,
"preview": "#pragma once\n#include <torch/extension.h>\n#include <THC/THC.h>\n\nvoid knn_device(float* ref_dev, int ref_width,\n float"
},
{
"path": "knn/src/knn.h",
"chars": 1795,
"preview": "#pragma once\n#include \"cpu/vision.h\"\n\n#ifdef WITH_CUDA\n#include \"cuda/vision.h\"\n#include <THC/THC.h>\nextern THCState *st"
},
{
"path": "knn/src/vision.cpp",
"chars": 108,
"preview": "#include \"knn.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"knn\", &knn, \"k-nearest neighbors\");\n}\n"
},
{
"path": "models/backbone_resunet14.py",
"chars": 7039,
"preview": "import MinkowskiEngine as ME\nfrom MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck\nfrom models.resnet "
},
{
"path": "models/graspnet.py",
"chars": 5941,
"preview": "\"\"\" GraspNet baseline model definition.\n Author: chenxi-wang\n\"\"\"\n\nimport os\nimport sys\nimport numpy as np\nimport torc"
},
{
"path": "models/loss.py",
"chars": 3361,
"preview": "import torch.nn as nn\nimport torch\n\n\ndef get_loss(end_points):\n objectness_loss, end_points = compute_objectness_loss"
},
{
"path": "models/modules.py",
"chars": 5421,
"preview": "import os\nimport sys\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nBASE_DIR = os.path.dirname(os.p"
},
{
"path": "models/resnet.py",
"chars": 5782,
"preview": "import torch.nn as nn\n\ntry:\n import open3d as o3d\nexcept ImportError:\n raise ImportError(\"Please install open3d wi"
},
{
"path": "pointnet2/_ext_src/include/ball_query.h",
"chars": 346,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/include/cuda_utils.h",
"chars": 1486,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/include/cylinder_query.h",
"chars": 243,
"preview": "// Author: chenxi-wang\n\n#pragma once\n#include <torch/extension.h>\n\nat::Tensor cylinder_query(at::Tensor new_xyz, at::Ten"
},
{
"path": "pointnet2/_ext_src/include/group_points.h",
"chars": 366,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/include/interpolate.h",
"chars": 569,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/include/sampling.h",
"chars": 443,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/include/utils.h",
"chars": 1185,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/ball_query.cpp",
"chars": 1224,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/ball_query_gpu.cu",
"chars": 1967,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/bindings.cpp",
"chars": 826,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/cylinder_query.cpp",
"chars": 1275,
"preview": "// Author: chenxi-wang\n\n#include \"cylinder_query.h\"\n#include \"utils.h\"\n\nvoid query_cylinder_point_kernel_wrapper(int b, "
},
{
"path": "pointnet2/_ext_src/src/cylinder_query_gpu.cu",
"chars": 2338,
"preview": "// Author: chenxi-wang\n\n#include <math.h>\n#include <stdio.h>\n#include <stdlib.h>\n\n#include \"cuda_utils.h\"\n\n__global__ vo"
},
{
"path": "pointnet2/_ext_src/src/group_points.cpp",
"chars": 2103,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/group_points_gpu.cu",
"chars": 3068,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/interpolate.cpp",
"chars": 3487,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/interpolate_gpu.cu",
"chars": 5324,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/sampling.cpp",
"chars": 3044,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/_ext_src/src/sampling_gpu.cu",
"chars": 7202,
"preview": "// Copyright (c) Facebook, Inc. and its affiliates.\n// \n// This source code is licensed under the MIT license found in t"
},
{
"path": "pointnet2/pointnet2_modules.py",
"chars": 17609,
"preview": "# Copyright (c) Facebook, Inc. and its affiliates.\n# \n# This source code is licensed under the MIT license found in the\n"
},
{
"path": "pointnet2/pointnet2_utils.py",
"chars": 16909,
"preview": "# Copyright (c) Facebook, Inc. and its affiliates.\n# \n# This source code is licensed under the MIT license found in the\n"
},
{
"path": "pointnet2/pytorch_utils.py",
"chars": 7501,
"preview": "# Copyright (c) Facebook, Inc. and its affiliates.\n# \n# This source code is licensed under the MIT license found in the\n"
},
{
"path": "pointnet2/setup.py",
"chars": 1006,
"preview": "# Copyright (c) Facebook, Inc. and its affiliates.\n# \n# This source code is licensed under the MIT license found in the\n"
},
{
"path": "requirements.txt",
"chars": 86,
"preview": "torch>=1.8\ntensorboard==2.3\nnumpy\nscipy\nopen3d>=0.8\nPillow\ntqdm\nMinkowskiEngine==0.5.4"
},
{
"path": "test.py",
"chars": 5357,
"preview": "import os\nimport sys\nimport numpy as np\nimport argparse\nimport time\nimport torch\nfrom torch.utils.data import DataLoader"
},
{
"path": "train.py",
"chars": 6300,
"preview": "import os\nimport sys\nimport numpy as np\nfrom datetime import datetime\nimport argparse\n\nimport torch\nimport torch.optim a"
},
{
"path": "utils/collision_detector.py",
"chars": 6596,
"preview": "\"\"\" Collision detection to remove collided grasp pose predictions.\nAuthor: chenxi-wang\n\"\"\"\n\nimport os\nimport sys\nimport "
},
{
"path": "utils/data_utils.py",
"chars": 5795,
"preview": "\"\"\" Tools for data processing.\n Author: chenxi-wang\n\"\"\"\n\nimport numpy as np\n\n\nclass CameraInfo():\n \"\"\" Camera intr"
},
{
"path": "utils/label_generation.py",
"chars": 7666,
"preview": "\"\"\" Dynamically generate grasp labels during training.\n Author: chenxi-wang\n\"\"\"\n\nimport os\nimport sys\nimport torch\n\nB"
},
{
"path": "utils/loss_utils.py",
"chars": 4404,
"preview": "\"\"\" Tools for loss computation.\n Author: chenxi-wang\n\"\"\"\n\nimport torch\nimport numpy as np\n\nGRASP_MAX_WIDTH = 0.1\nGRAS"
}
]
// ... and 1 more files (download for full content)
About this extraction
This page contains the full source code of the rhett-chen/graspness_implementation GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 53 files (198.5 KB), approximately 53.3k tokens, and a symbol index with 195 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.