├── doc ├── first.png └── example_data │ ├── color.png │ ├── depth.png │ ├── meta.mat │ ├── demo_result.png │ ├── orange_color.png │ ├── orange_depth.raw │ └── workspace_mask.png ├── requirements.txt ├── models ├── __pycache__ │ ├── loss.cpython-36.pyc │ ├── loss.cpython-37.pyc │ ├── decode.cpython-36.pyc │ ├── modules.cpython-36.pyc │ ├── modules.cpython-37.pyc │ ├── backbone.cpython-36.pyc │ ├── backbone.cpython-37.pyc │ ├── graspnet.cpython-36.pyc │ ├── graspnet.cpython-37.pyc │ ├── no_local.cpython-36.pyc │ ├── nocrop_net.cpython-36.pyc │ ├── FGC_graspnet.cpython-36.pyc │ ├── ablation_net.cpython-36.pyc │ ├── best_nolocal.cpython-36.pyc │ ├── dual_attention.cpython-36.pyc │ ├── single_graspnet.cpython-36.pyc │ └── multi_scale_backbone.cpython-36.pyc ├── decode.py ├── backbone.py ├── FGC_graspnet.py ├── modules.py └── loss.py ├── utils ├── __pycache__ │ ├── misc.cpython-36.pyc │ ├── data_utils.cpython-36.pyc │ ├── data_utils.cpython-37.pyc │ ├── loss_utils.cpython-36.pyc │ ├── loss_utils.cpython-37.pyc │ ├── label_generation.cpython-36.pyc │ ├── label_generation.cpython-37.pyc │ ├── collision_detector.cpython-36.pyc │ └── collision_detector.cpython-37.pyc ├── loss_utils.py ├── data_utils.py ├── collision_detector.py └── label_generation.py ├── knn ├── __pycache__ │ └── knn_modules.cpython-36.pyc ├── src │ ├── vision.cpp │ ├── cpu │ │ ├── vision.h │ │ └── knn_cpu.cpp │ ├── cuda │ │ ├── vision.h │ │ └── knn.cu │ └── knn.h ├── knn_modules.py └── setup.py ├── FGC_generate ├── __pycache__ │ ├── vector.cpython-36.pyc │ ├── vis_grasp.cpython-36.pyc │ └── contact_score.cpython-36.pyc ├── README.md ├── score_allocation.py ├── vector.py ├── rnn_neighbor.py ├── vis_grasp.py ├── contact_score.py ├── contact_score_numba.py └── score_gen_mp.py ├── dataset └── __pycache__ │ ├── graspnet_dataset.cpython-36.pyc │ ├── graspnet_dataset.cpython-37.pyc │ └── graspnet_dataset_init.cpython-36.pyc ├── pointnet2 ├── __pycache__ │ ├── pytorch_utils.cpython-36.pyc │ ├── pointnet2_modules.cpython-36.pyc │ └── pointnet2_utils.cpython-36.pyc ├── _ext_src │ ├── include │ │ ├── cylinder_query.h │ │ ├── ball_query.h │ │ ├── group_points.h │ │ ├── sampling.h │ │ ├── interpolate.h │ │ ├── utils.h │ │ └── cuda_utils.h │ └── src │ │ ├── bindings.cpp │ │ ├── ball_query.cpp │ │ ├── cylinder_query.cpp │ │ ├── ball_query_gpu.cu │ │ ├── group_points.cpp │ │ ├── cylinder_query_gpu.cu │ │ ├── group_points_gpu.cu │ │ ├── sampling.cpp │ │ ├── interpolate.cpp │ │ ├── interpolate_gpu.cu │ │ └── sampling_gpu.cu ├── setup.py └── pytorch_utils.py ├── test.sh ├── train.sh ├── README.md ├── demo.py ├── test.py └── log └── log_train.txt /doc/first.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/doc/first.png -------------------------------------------------------------------------------- /doc/example_data/color.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/doc/example_data/color.png -------------------------------------------------------------------------------- /doc/example_data/depth.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/doc/example_data/depth.png -------------------------------------------------------------------------------- /doc/example_data/meta.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/doc/example_data/meta.mat -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch==1.6 2 | tensorboard==2.3 3 | numpy 4 | scipy 5 | open3d==0.8 6 | Pillow 7 | tqdm 8 | -------------------------------------------------------------------------------- /doc/example_data/demo_result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/doc/example_data/demo_result.png -------------------------------------------------------------------------------- /doc/example_data/orange_color.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/doc/example_data/orange_color.png -------------------------------------------------------------------------------- /doc/example_data/orange_depth.raw: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/doc/example_data/orange_depth.raw -------------------------------------------------------------------------------- /doc/example_data/workspace_mask.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/doc/example_data/workspace_mask.png -------------------------------------------------------------------------------- /models/__pycache__/loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/loss.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/loss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/loss.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/misc.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/misc.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/decode.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/decode.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/modules.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/modules.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/modules.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/modules.cpython-37.pyc -------------------------------------------------------------------------------- /knn/__pycache__/knn_modules.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/knn/__pycache__/knn_modules.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/backbone.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/backbone.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/backbone.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/backbone.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/graspnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/graspnet.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/graspnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/graspnet.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/no_local.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/no_local.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/nocrop_net.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/nocrop_net.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/data_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/data_utils.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/data_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/data_utils.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/loss_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/loss_utils.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/loss_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/loss_utils.cpython-37.pyc -------------------------------------------------------------------------------- /FGC_generate/__pycache__/vector.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/FGC_generate/__pycache__/vector.cpython-36.pyc -------------------------------------------------------------------------------- /knn/src/vision.cpp: -------------------------------------------------------------------------------- 1 | #include "knn.h" 2 | 3 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 4 | m.def("knn", &knn, "k-nearest neighbors"); 5 | } 6 | -------------------------------------------------------------------------------- /models/__pycache__/FGC_graspnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/FGC_graspnet.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/ablation_net.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/ablation_net.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/best_nolocal.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/best_nolocal.cpython-36.pyc -------------------------------------------------------------------------------- /FGC_generate/__pycache__/vis_grasp.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/FGC_generate/__pycache__/vis_grasp.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/dual_attention.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/dual_attention.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/single_graspnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/single_graspnet.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/label_generation.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/label_generation.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/label_generation.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/label_generation.cpython-37.pyc -------------------------------------------------------------------------------- /dataset/__pycache__/graspnet_dataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/dataset/__pycache__/graspnet_dataset.cpython-36.pyc -------------------------------------------------------------------------------- /dataset/__pycache__/graspnet_dataset.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/dataset/__pycache__/graspnet_dataset.cpython-37.pyc -------------------------------------------------------------------------------- /pointnet2/__pycache__/pytorch_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/pointnet2/__pycache__/pytorch_utils.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/collision_detector.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/collision_detector.cpython-36.pyc -------------------------------------------------------------------------------- /utils/__pycache__/collision_detector.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/utils/__pycache__/collision_detector.cpython-37.pyc -------------------------------------------------------------------------------- /FGC_generate/__pycache__/contact_score.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/FGC_generate/__pycache__/contact_score.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/multi_scale_backbone.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/models/__pycache__/multi_scale_backbone.cpython-36.pyc -------------------------------------------------------------------------------- /pointnet2/__pycache__/pointnet2_modules.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/pointnet2/__pycache__/pointnet2_modules.cpython-36.pyc -------------------------------------------------------------------------------- /pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/pointnet2/__pycache__/pointnet2_utils.cpython-36.pyc -------------------------------------------------------------------------------- /dataset/__pycache__/graspnet_dataset_init.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/luyh20/FGC-GraspNet/HEAD/dataset/__pycache__/graspnet_dataset_init.cpython-36.pyc -------------------------------------------------------------------------------- /test.sh: -------------------------------------------------------------------------------- 1 | export CUDA_VISIBLE_DEVICES=0 2 | python test.py --checkpoint_path logs/baseline/checkpoint.tar --camera realsense --dataset_root ./grasp_data --batch_size 3 3 | -------------------------------------------------------------------------------- /train.sh: -------------------------------------------------------------------------------- 1 | export CUDA_VISIBLE_DEVICES=4,5,6,7 2 | python -m torch.distributed.launch --nproc_per_node=4 --use_env train.py --camera realsense --log_dir logs/baseline --batch_size 3 --dataset_root grasp_data/ -------------------------------------------------------------------------------- /knn/src/cpu/vision.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | void knn_cpu(float* ref_dev, int ref_width, 5 | float* query_dev, int query_width, 6 | int height, int k, float* dist_dev, long* ind_dev, long* ind_buf); -------------------------------------------------------------------------------- /knn/src/cuda/vision.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | 5 | void knn_device(float* ref_dev, int ref_width, 6 | float* query_dev, int query_width, 7 | int height, int k, float* dist_dev, long* ind_dev, cudaStream_t stream); -------------------------------------------------------------------------------- /FGC_generate/README.md: -------------------------------------------------------------------------------- 1 | ## Introduction 2 | 3 | The **FGC_generate** is about how to generate the new grasp score labels by hybrid physical metric. 4 | This code is only for reference if you want to know the process of FGC scores generation, and is not related to model training or testing. 5 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/cylinder_query.h: -------------------------------------------------------------------------------- 1 | // Author: chenxi-wang 2 | 3 | #pragma once 4 | #include 5 | 6 | at::Tensor cylinder_query(at::Tensor new_xyz, at::Tensor xyz, at::Tensor rot, const float radius, const float hmin, const float hmax, 7 | const int nsample); 8 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/ball_query.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 10 | const int nsample); 11 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/group_points.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor group_points(at::Tensor points, at::Tensor idx); 10 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 11 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/sampling.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | 9 | at::Tensor gather_points(at::Tensor points, at::Tensor idx); 10 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n); 11 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples); 12 | -------------------------------------------------------------------------------- /knn/knn_modules.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import gc 3 | import operator as op 4 | import functools 5 | import torch 6 | from torch.autograd import Variable, Function 7 | from knn_pytorch import knn_pytorch 8 | # import knn_pytorch 9 | def knn(ref, query, k=1): 10 | """ Compute k nearest neighbors for each query point. 11 | """ 12 | device = ref.device 13 | ref = ref.float().to(device) 14 | query = query.float().to(device) 15 | inds = torch.empty(query.shape[0], k, query.shape[2]).long().to(device) 16 | knn_pytorch.knn(ref, query, inds) 17 | return inds 18 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/interpolate.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | 8 | #include 9 | #include 10 | 11 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows); 12 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 13 | at::Tensor weight); 14 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 15 | at::Tensor weight, const int m); 16 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/bindings.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "ball_query.h" 7 | #include "group_points.h" 8 | #include "interpolate.h" 9 | #include "sampling.h" 10 | #include "cylinder_query.h" 11 | 12 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 13 | m.def("gather_points", &gather_points); 14 | m.def("gather_points_grad", &gather_points_grad); 15 | m.def("furthest_point_sampling", &furthest_point_sampling); 16 | 17 | m.def("three_nn", &three_nn); 18 | m.def("three_interpolate", &three_interpolate); 19 | m.def("three_interpolate_grad", &three_interpolate_grad); 20 | 21 | m.def("ball_query", &ball_query); 22 | 23 | m.def("group_points", &group_points); 24 | m.def("group_points_grad", &group_points_grad); 25 | 26 | m.def("cylinder_query", &cylinder_query); 27 | } 28 | -------------------------------------------------------------------------------- /pointnet2/setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | from setuptools import setup 7 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 8 | import glob 9 | import os 10 | ROOT = os.path.dirname(os.path.abspath(__file__)) 11 | 12 | _ext_src_root = "_ext_src" 13 | _ext_sources = glob.glob("{}/src/*.cpp".format(_ext_src_root)) + glob.glob( 14 | "{}/src/*.cu".format(_ext_src_root) 15 | ) 16 | _ext_headers = glob.glob("{}/include/*".format(_ext_src_root)) 17 | 18 | setup( 19 | name='pointnet2', 20 | ext_modules=[ 21 | CUDAExtension( 22 | name='pointnet2._ext', 23 | sources=_ext_sources, 24 | extra_compile_args={ 25 | "cxx": ["-O2", "-I{}".format("{}/{}/include".format(ROOT, _ext_src_root))], 26 | "nvcc": ["-O2", "-I{}".format("{}/{}/include".format(ROOT, _ext_src_root))], 27 | }, 28 | ) 29 | ], 30 | cmdclass={ 31 | 'build_ext': BuildExtension 32 | } 33 | ) 34 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #pragma once 7 | #include 8 | #include 9 | 10 | #define CHECK_CUDA(x) \ 11 | do { \ 12 | TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \ 13 | } while (0) 14 | 15 | #define CHECK_CONTIGUOUS(x) \ 16 | do { \ 17 | TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \ 18 | } while (0) 19 | 20 | #define CHECK_IS_INT(x) \ 21 | do { \ 22 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \ 23 | #x " must be an int tensor"); \ 24 | } while (0) 25 | 26 | #define CHECK_IS_FLOAT(x) \ 27 | do { \ 28 | TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \ 29 | #x " must be a float tensor"); \ 30 | } while (0) 31 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/ball_query.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "ball_query.h" 7 | #include "utils.h" 8 | 9 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 10 | int nsample, const float *new_xyz, 11 | const float *xyz, int *idx); 12 | 13 | at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, 14 | const int nsample) { 15 | CHECK_CONTIGUOUS(new_xyz); 16 | CHECK_CONTIGUOUS(xyz); 17 | CHECK_IS_FLOAT(new_xyz); 18 | CHECK_IS_FLOAT(xyz); 19 | 20 | if (new_xyz.type().is_cuda()) { 21 | CHECK_CUDA(xyz); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 26 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 27 | 28 | if (new_xyz.type().is_cuda()) { 29 | query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 30 | radius, nsample, new_xyz.data(), 31 | xyz.data(), idx.data()); 32 | } else { 33 | TORCH_CHECK(false, "CPU not supported"); 34 | } 35 | 36 | return idx; 37 | } 38 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/cylinder_query.cpp: -------------------------------------------------------------------------------- 1 | // Author: chenxi-wang 2 | 3 | #include "cylinder_query.h" 4 | #include "utils.h" 5 | 6 | void query_cylinder_point_kernel_wrapper(int b, int n, int m, float radius, float hmin, float hmax, 7 | int nsample, const float *new_xyz, 8 | const float *xyz, const float *rot, int *idx); 9 | 10 | at::Tensor cylinder_query(at::Tensor new_xyz, at::Tensor xyz, at::Tensor rot, const float radius, const float hmin, const float hmax, 11 | const int nsample) { 12 | CHECK_CONTIGUOUS(new_xyz); 13 | CHECK_CONTIGUOUS(xyz); 14 | CHECK_CONTIGUOUS(rot); 15 | CHECK_IS_FLOAT(new_xyz); 16 | CHECK_IS_FLOAT(xyz); 17 | CHECK_IS_FLOAT(rot); 18 | 19 | if (new_xyz.type().is_cuda()) { 20 | CHECK_CUDA(xyz); 21 | CHECK_CUDA(rot); 22 | } 23 | 24 | at::Tensor idx = 25 | torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample}, 26 | at::device(new_xyz.device()).dtype(at::ScalarType::Int)); 27 | 28 | if (new_xyz.type().is_cuda()) { 29 | query_cylinder_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1), 30 | radius, hmin, hmax, nsample, new_xyz.data(), 31 | xyz.data(), rot.data(), idx.data()); 32 | } else { 33 | TORCH_CHECK(false, "CPU not supported"); 34 | } 35 | 36 | return idx; 37 | } 38 | -------------------------------------------------------------------------------- /FGC_generate/score_allocation.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | 4 | 5 | root = '../grasp_data' 6 | 7 | for i in range(87): 8 | eg = np.load(os.path.join(root, 'new_score', '{}_labels.npz'.format(str(i).zfill(3)))) 9 | s5 = np.load(os.path.join(root, 'score5', '{}_labels.npz'.format(str(i).zfill(3)))) 10 | print(eg.files) 11 | mask_idx = eg['mask_idx'] 12 | score2 = eg['score2'] 13 | score3 = eg['score3'] 14 | score4 = eg['score4'] 15 | score5 = s5['score5'] 16 | 17 | score2_norm = 1 - (score2-np.min(score2))/(np.max(score2)-np.min(score2)) 18 | 19 | obj_path = os.path.join(root, 'grasp_label', '{}_labels.npz'.format(str(i).zfill(3))) 20 | label = np.load(obj_path) 21 | label_score = label['scores'] 22 | mask1 = (label_score > 0) 23 | fric_coefs = label_score[mask1] 24 | scores = 1.1 - fric_coefs 25 | 26 | idx = np.array([220872, 434776, 213653, 8151, 222708]) 27 | 28 | final_score = 0.7*scores+0.05*score2_norm+0.2*score3*score4+0.05*score5 29 | final_score_norm = (final_score-np.min(final_score))/(np.max(final_score)-np.min(final_score)) 30 | # idx = np.array([220872, 434776, 213653, 8151, 222708]) 31 | print(score2_norm[idx]) 32 | print((score3 * score4)[idx]) 33 | # print(0.05*score5[idx]) 34 | # print(final_score[idx]) 35 | # print(final_score_norm[idx]) 36 | label_score[mask1] = final_score_norm 37 | savepath = os.path.join(root, 'new_grasp_label_2345', '{}_labels.npz'.format(str(i).zfill(3))) 38 | np.savez(savepath, new_scores=label_score) 39 | 40 | print('1') -------------------------------------------------------------------------------- /pointnet2/_ext_src/include/cuda_utils.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #ifndef _CUDA_UTILS_H 7 | #define _CUDA_UTILS_H 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | 16 | #include 17 | 18 | #define TOTAL_THREADS 512 19 | 20 | inline int opt_n_threads(int work_size) { 21 | const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 22 | 23 | return max(min(1 << pow_2, TOTAL_THREADS), 1); 24 | } 25 | 26 | inline dim3 opt_block_config(int x, int y) { 27 | const int x_threads = opt_n_threads(x); 28 | const int y_threads = 29 | max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 30 | dim3 block_config(x_threads, y_threads, 1); 31 | 32 | return block_config; 33 | } 34 | 35 | #define CUDA_CHECK_ERRORS() \ 36 | do { \ 37 | cudaError_t err = cudaGetLastError(); \ 38 | if (cudaSuccess != err) { \ 39 | fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \ 40 | cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \ 41 | __FILE__); \ 42 | exit(-1); \ 43 | } \ 44 | } while (0) 45 | 46 | #endif 47 | -------------------------------------------------------------------------------- /knn/src/cpu/knn_cpu.cpp: -------------------------------------------------------------------------------- 1 | #include "cpu/vision.h" 2 | 3 | 4 | void knn_cpu(float* ref_dev, int ref_width, float* query_dev, int query_width, 5 | int height, int k, float* dist_dev, long* ind_dev, long* ind_buf) 6 | { 7 | // Compute all the distances 8 | for(int query_idx = 0;query_idx dist_dev[query_idx * ref_width + j + 1]) 31 | { 32 | temp_value = dist_dev[query_idx * ref_width + j]; 33 | dist_dev[query_idx * ref_width + j] = dist_dev[query_idx * ref_width + j + 1]; 34 | dist_dev[query_idx * ref_width + j + 1] = temp_value; 35 | temp_idx = ind_buf[j]; 36 | ind_buf[j] = ind_buf[j + 1]; 37 | ind_buf[j + 1] = temp_idx; 38 | } 39 | 40 | } 41 | 42 | for(int i = 0;i < k;i++) 43 | ind_dev[query_idx + i * query_width] = ind_buf[i]; 44 | #if DEBUG 45 | for(int i = 0;i < ref_width;i++) 46 | printf("%d, ", ind_buf[i]); 47 | printf("\n"); 48 | #endif 49 | 50 | } 51 | 52 | 53 | 54 | 55 | 56 | } -------------------------------------------------------------------------------- /knn/src/knn.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "cpu/vision.h" 3 | 4 | #ifdef WITH_CUDA 5 | #include "cuda/vision.h" 6 | #include 7 | extern THCState *state; 8 | #endif 9 | 10 | 11 | 12 | int knn(at::Tensor& ref, at::Tensor& query, at::Tensor& idx) 13 | { 14 | 15 | // TODO check dimensions 16 | long batch, ref_nb, query_nb, dim, k; 17 | batch = ref.size(0); 18 | dim = ref.size(1); 19 | k = idx.size(1); 20 | ref_nb = ref.size(2); 21 | query_nb = query.size(2); 22 | 23 | float *ref_dev = ref.data(); 24 | float *query_dev = query.data(); 25 | long *idx_dev = idx.data(); 26 | 27 | 28 | 29 | 30 | if (ref.type().is_cuda()) { 31 | #ifdef WITH_CUDA 32 | // TODO raise error if not compiled with CUDA 33 | float *dist_dev = (float*)THCudaMalloc(state, ref_nb * query_nb * sizeof(float)); 34 | 35 | for (int b = 0; b < batch; b++) 36 | { 37 | // knn_device(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k, 38 | // dist_dev, idx_dev + b * k * query_nb, THCState_getCurrentStream(state)); 39 | knn_device(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k, 40 | dist_dev, idx_dev + b * k * query_nb, c10::cuda::getCurrentCUDAStream()); 41 | } 42 | THCudaFree(state, dist_dev); 43 | cudaError_t err = cudaGetLastError(); 44 | if (err != cudaSuccess) 45 | { 46 | printf("error in knn: %s\n", cudaGetErrorString(err)); 47 | THError("aborting"); 48 | } 49 | return 1; 50 | #else 51 | AT_ERROR("Not compiled with GPU support"); 52 | #endif 53 | } 54 | 55 | 56 | float *dist_dev = (float*)malloc(ref_nb * query_nb * sizeof(float)); 57 | long *ind_buf = (long*)malloc(ref_nb * sizeof(long)); 58 | for (int b = 0; b < batch; b++) { 59 | knn_cpu(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k, 60 | dist_dev, idx_dev + b * k * query_nb, ind_buf); 61 | } 62 | 63 | free(dist_dev); 64 | free(ind_buf); 65 | 66 | return 1; 67 | 68 | } 69 | -------------------------------------------------------------------------------- /knn/setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import glob 4 | import os 5 | 6 | import torch 7 | from setuptools import find_packages 8 | from setuptools import setup 9 | from torch.utils.cpp_extension import CUDA_HOME 10 | from torch.utils.cpp_extension import CppExtension 11 | from torch.utils.cpp_extension import CUDAExtension 12 | 13 | requirements = ["torch", "torchvision"] 14 | 15 | 16 | def get_extensions(): 17 | this_dir = os.path.dirname(os.path.abspath(__file__)) 18 | extensions_dir = os.path.join(this_dir, "src") 19 | 20 | main_file = glob.glob(os.path.join(extensions_dir, "*.cpp")) 21 | source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp")) 22 | source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu")) 23 | 24 | sources = main_file + source_cpu 25 | extension = CppExtension 26 | 27 | extra_compile_args = {"cxx": []} 28 | define_macros = [] 29 | 30 | if torch.cuda.is_available() and CUDA_HOME is not None: 31 | extension = CUDAExtension 32 | sources += source_cuda 33 | define_macros += [("WITH_CUDA", None)] 34 | extra_compile_args["nvcc"] = [ 35 | "-DCUDA_HAS_FP16=1", 36 | "-D__CUDA_NO_HALF_OPERATORS__", 37 | "-D__CUDA_NO_HALF_CONVERSIONS__", 38 | "-D__CUDA_NO_HALF2_OPERATORS__", 39 | ] 40 | 41 | sources = [os.path.join(extensions_dir, s) for s in sources] 42 | 43 | include_dirs = [extensions_dir] 44 | 45 | ext_modules = [ 46 | extension( 47 | "knn_pytorch.knn_pytorch", 48 | sources, 49 | include_dirs=include_dirs, 50 | define_macros=define_macros, 51 | extra_compile_args=extra_compile_args, 52 | ) 53 | ] 54 | 55 | return ext_modules 56 | 57 | 58 | setup( 59 | name="knn_pytorch", 60 | version="0.1", 61 | author="foolyc", 62 | url="https://github.com/foolyc/torchKNN", 63 | description="KNN implement in Pytorch 1.0 including both cpu version and gpu version", 64 | ext_modules=get_extensions(), 65 | cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}, 66 | ) 67 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/ball_query_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | #include 9 | 10 | #include "cuda_utils.h" 11 | 12 | // input: new_xyz(b, m, 3) xyz(b, n, 3) 13 | // output: idx(b, m, nsample) 14 | __global__ void query_ball_point_kernel(int b, int n, int m, float radius, 15 | int nsample, 16 | const float *__restrict__ new_xyz, 17 | const float *__restrict__ xyz, 18 | int *__restrict__ idx) { 19 | int batch_index = blockIdx.x; 20 | xyz += batch_index * n * 3; 21 | new_xyz += batch_index * m * 3; 22 | idx += m * nsample * batch_index; 23 | 24 | int index = threadIdx.x; 25 | int stride = blockDim.x; 26 | 27 | float radius2 = radius * radius; 28 | for (int j = index; j < m; j += stride) { 29 | float new_x = new_xyz[j * 3 + 0]; 30 | float new_y = new_xyz[j * 3 + 1]; 31 | float new_z = new_xyz[j * 3 + 2]; 32 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 33 | float x = xyz[k * 3 + 0]; 34 | float y = xyz[k * 3 + 1]; 35 | float z = xyz[k * 3 + 2]; 36 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + 37 | (new_z - z) * (new_z - z); 38 | if (d2 < radius2) { 39 | if (cnt == 0) { 40 | for (int l = 0; l < nsample; ++l) { 41 | idx[j * nsample + l] = k; 42 | } 43 | } 44 | idx[j * nsample + cnt] = k; 45 | ++cnt; 46 | } 47 | } 48 | } 49 | } 50 | 51 | void query_ball_point_kernel_wrapper(int b, int n, int m, float radius, 52 | int nsample, const float *new_xyz, 53 | const float *xyz, int *idx) { 54 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 55 | query_ball_point_kernel<<>>( 56 | b, n, m, radius, nsample, new_xyz, xyz, idx); 57 | 58 | CUDA_CHECK_ERRORS(); 59 | } 60 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/group_points.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "group_points.h" 7 | #include "utils.h" 8 | 9 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 10 | const float *points, const int *idx, 11 | float *out); 12 | 13 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 14 | int nsample, const float *grad_out, 15 | const int *idx, float *grad_points); 16 | 17 | at::Tensor group_points(at::Tensor points, at::Tensor idx) { 18 | CHECK_CONTIGUOUS(points); 19 | CHECK_CONTIGUOUS(idx); 20 | CHECK_IS_FLOAT(points); 21 | CHECK_IS_INT(idx); 22 | 23 | if (points.type().is_cuda()) { 24 | CHECK_CUDA(idx); 25 | } 26 | 27 | at::Tensor output = 28 | torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)}, 29 | at::device(points.device()).dtype(at::ScalarType::Float)); 30 | 31 | if (points.type().is_cuda()) { 32 | group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 33 | idx.size(1), idx.size(2), points.data(), 34 | idx.data(), output.data()); 35 | } else { 36 | TORCH_CHECK(false, "CPU not supported"); 37 | } 38 | 39 | return output; 40 | } 41 | 42 | at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) { 43 | CHECK_CONTIGUOUS(grad_out); 44 | CHECK_CONTIGUOUS(idx); 45 | CHECK_IS_FLOAT(grad_out); 46 | CHECK_IS_INT(idx); 47 | 48 | if (grad_out.type().is_cuda()) { 49 | CHECK_CUDA(idx); 50 | } 51 | 52 | at::Tensor output = 53 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 54 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 55 | 56 | if (grad_out.type().is_cuda()) { 57 | group_points_grad_kernel_wrapper( 58 | grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2), 59 | grad_out.data(), idx.data(), output.data()); 60 | } else { 61 | TORCH_CHECK(false, "CPU not supported"); 62 | } 63 | 64 | return output; 65 | } 66 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/cylinder_query_gpu.cu: -------------------------------------------------------------------------------- 1 | // Author: chenxi-wang 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | #include "cuda_utils.h" 8 | 9 | __global__ void query_cylinder_point_kernel(int b, int n, int m, float radius, float hmin, float hmax, 10 | int nsample, 11 | const float *__restrict__ new_xyz, 12 | const float *__restrict__ xyz, 13 | const float *__restrict__ rot, 14 | int *__restrict__ idx) { 15 | int batch_index = blockIdx.x; 16 | xyz += batch_index * n * 3; 17 | new_xyz += batch_index * m * 3; 18 | rot += batch_index * m * 9; 19 | idx += m * nsample * batch_index; 20 | 21 | int index = threadIdx.x; 22 | int stride = blockDim.x; 23 | 24 | float radius2 = radius * radius; 25 | for (int j = index; j < m; j += stride) { 26 | float new_x = new_xyz[j * 3 + 0]; 27 | float new_y = new_xyz[j * 3 + 1]; 28 | float new_z = new_xyz[j * 3 + 2]; 29 | float r0 = rot[j * 9 + 0]; 30 | float r1 = rot[j * 9 + 1]; 31 | float r2 = rot[j * 9 + 2]; 32 | float r3 = rot[j * 9 + 3]; 33 | float r4 = rot[j * 9 + 4]; 34 | float r5 = rot[j * 9 + 5]; 35 | float r6 = rot[j * 9 + 6]; 36 | float r7 = rot[j * 9 + 7]; 37 | float r8 = rot[j * 9 + 8]; 38 | for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) { 39 | float x = xyz[k * 3 + 0] - new_x; 40 | float y = xyz[k * 3 + 1] - new_y; 41 | float z = xyz[k * 3 + 2] - new_z; 42 | float x_rot = r0 * x + r3 * y + r6 * z; 43 | float y_rot = r1 * x + r4 * y + r7 * z; 44 | float z_rot = r2 * x + r5 * y + r8 * z; 45 | float d2 = y_rot * y_rot + z_rot * z_rot; 46 | if (d2 < radius2 && x_rot > hmin && x_rot < hmax) { 47 | if (cnt == 0) { 48 | for (int l = 0; l < nsample; ++l) { 49 | idx[j * nsample + l] = k; 50 | } 51 | } 52 | idx[j * nsample + cnt] = k; 53 | ++cnt; 54 | } 55 | } 56 | } 57 | } 58 | 59 | void query_cylinder_point_kernel_wrapper(int b, int n, int m, float radius, float hmin, float hmax, 60 | int nsample, const float *new_xyz, 61 | const float *xyz, const float *rot, int *idx) { 62 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 63 | query_cylinder_point_kernel<<>>( 64 | b, n, m, radius, hmin, hmax, nsample, new_xyz, xyz, rot, idx); 65 | 66 | CUDA_CHECK_ERRORS(); 67 | } 68 | -------------------------------------------------------------------------------- /FGC_generate/vector.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d 2 | import numpy as np 3 | 4 | 5 | def PCA(data, correlation=False, sort=True): 6 | 7 | average_data = np.mean(data,axis=0) #求 NX3 向量的均值 8 | decentration_matrix = data - average_data #去中心化 9 | H = np.dot(decentration_matrix.T,decentration_matrix) #求解协方差矩阵 H 10 | eigenvectors,eigenvalues,eigenvectors_T = np.linalg.svd(H) # SVD求解特征值、特征向量 11 | 12 | 13 | if sort: 14 | sort = eigenvalues.argsort()[::-1] #降序排列 15 | eigenvalues = eigenvalues[sort] #索引 16 | eigenvectors = eigenvectors[:, sort] 17 | 18 | return eigenvalues, eigenvectors 19 | 20 | 21 | def get_vec_pca(pco3d): 22 | pcd_tree = o3d.geometry.KDTreeFlann(pco3d) 23 | normals = [] 24 | 25 | pc = np.asarray(pco3d.points) 26 | print(pc.shape[0]) 27 | for i in range(pc.shape[0]): 28 | [_,idx,_] = pcd_tree.search_knn_vector_3d(pco3d.points[i], 10) #取10个临近点进行曲线拟合 29 | # asarray和array 一样 但是array会copy出一个副本，asarray不会，节省内存 30 | k_nearest_point = np.asarray(pco3d.points)[idx, :] # 找出每一点的10个临近点，类似于拟合成曲面，然后进行PCA找到特征向量最小的值，作为法向量 31 | w, v = PCA(k_nearest_point) 32 | normals.append(v[:, 2]) 33 | 34 | normals = np.array(normals, dtype=np.float64) 35 | pco3d.normals = o3d.utility.Vector3dVector(normals) 36 | o3d.visualization.draw_geometries([pco3d], point_show_normal=True) 37 | print(np.asarray(pco3d.normals)) 38 | return normals 39 | 40 | 41 | def orient_normals(pcd): 42 | # 法向量方向一致性 43 | normals = np.asarray(pcd.normals) 44 | point = np.asarray(pcd.points) 45 | normals_new = np.zeros((normals.shape[0], normals.shape[1])) 46 | 47 | for i in range(point.shape[0]): 48 | vector_approach = point[i] 49 | normal = normals[i] 50 | orient = np.dot(vector_approach, normal.T) 51 | 52 | if orient < 0: # 法向量与中心向量相反，则法向量反向 53 | normals_new[i] = -normal 54 | else: 55 | normals_new[i] = normal 56 | return normals_new 57 | 58 | 59 | if __name__=='__main__': 60 | ob = np.load('067_labels.npz') 61 | print(ob.files) 62 | pc = ob['points'] 63 | 64 | pco3d = o3d.geometry.PointCloud() 65 | pco3d.points = o3d.utility.Vector3dVector(pc) 66 | pco3d.estimate_normals() 67 | a = np.asarray(pco3d.normals) 68 | pco3d.orient_normals_to_align_with_direction() 69 | 70 | o3d.visualization.draw_geometries([pco3d], "Open3D normal estimation", width=800, height=600, left=50, top=50, 71 | point_show_normal=True, mesh_show_wireframe=False, 72 | mesh_show_back_face=False) 73 | 74 | b = orient_normals(pco3d) 75 | pco3d.normals = o3d.utility.Vector3dVector(b) 76 | o3d.visualization.draw_geometries([pco3d], width=800, height=600, left=50, top=50, point_show_normal=True) 77 | c = a-b 78 | print(a) 79 | print(b) 80 | print(a-b) 81 | 82 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FGC-GraspNet 2 | Official Implementation for paper "Hybrid Physical Metric For 6-DoF Grasp Pose Detection" ICRA 2022. 3 | 4 | [arxiv](https://arxiv.org/abs/2206.11141) 5 | 6 | 7 | ![teaser](doc/first.png) 8 | 9 | ## Data Preparation 10 | Download the GraspNet-1Billion dataset from [graspnet](https://graspnet.net/datasets.html). 11 | In this paper, we use a new evaluation metric to generate grasp confidence scores for grasp poses. 12 | You can get the new score labels under hybrid physical metric from [here](https://drive.google.com/u/0/uc?id=1wAcGKOAO3EKWV0iih5sSVW5fHsu0X11R&export=download). 13 | The data directories should like this: 14 | 15 | ``` 16 | FGC_GraspNet/ 17 | ├── grasp_data/ 18 | | ├── scenes 19 | | ├── models 20 | | ├── dex_models 21 | │ ├── FGC_label 22 | │ ├── grasp_label 23 | │ └── collision_label 24 | ``` 25 | 26 | ## Requirements 27 | - Python 3.6.13 28 | - PyTorch 1.8.1+cu111 29 | - Open3d 0.8 30 | - TensorBoard 2.3 31 | - NumPy 32 | - SciPy 33 | - Pillow 34 | - tqdm 35 | 36 | ## Installation 37 | Get the code. 38 | ```bash 39 | git clone https://github.com/luyh20/FGC-GraspNet.git 40 | ``` 41 | Install packages via Pip. 42 | ```bash 43 | pip install -r requirements.txt 44 | ``` 45 | Compile and install pointnet2 operators (code adapted from [votenet](https://github.com/facebookresearch/votenet)). 46 | ```bash 47 | cd pointnet2 48 | python setup.py install 49 | ``` 50 | Compile and install knn operator (code adapted from [pytorch_knn_cuda](https://github.com/chrischoy/pytorch_knn_cuda)). 51 | ```bash 52 | cd knn 53 | python setup.py install 54 | ``` 55 | Install graspnetAPI. 56 | ```bash 57 | git clone https://github.com/graspnet/graspnetAPI.git 58 | cd graspnetAPI 59 | pip install . 60 | ``` 61 | 62 | 63 | ## Training and Testing 64 | Training: 65 | ``` 66 | sh train.sh 67 | ``` 68 | 69 | Testing: 70 | ``` 71 | sh test.sh 72 | ``` 73 | 74 | ## Model 75 | Realsense model's link，https://drive.google.com/file/d/1Y-CWHr_eZDoZm3XJocrUJq1SA5tfrONX/view?usp=sharing 76 | 77 | ## Demo 78 | The demo uses the RGBD data collected in real time from the Realsense D435i camera as input, and predicts the grasp poses results by the FGC_GraspNet. 79 | If you want to use it in your own experiment, you need to check the camera model or change the camera intrinsic for your camera model. 80 | 81 | Run Demo: 82 | ``` 83 | python demo.py 84 | ``` 85 | 86 | ## Video 87 | [![Video for FGC-GraspNet](https://i9.ytimg.com/vi_webp/gMoomsMJU_Y/mq1.webp?sqp=COSlpJsG-oaymwEmCMACELQB8quKqQMa8AEB-AHUBoAC3gOKAgwIABABGGUgYyhTMA8=&rs=AOn4CLAY_KqbACwxkemw7gMPLiYk8WuBkQ)](https://www.youtube.com/watch?v=gMoomsMJU_Y) 88 | 89 | ## BibTeX 90 | ``` 91 | @inproceedings{lu2022hybrid, 92 | title={Hybrid Physical Metric For 6-DoF Grasp Pose Detection}, 93 | author={Lu, Yuhao and Deng, Beixing and Wang, Zhenyu and Zhi, Peiyuan and Li, Yali and Wang, Shengjin}, 94 | booktitle={2022 International Conference on Robotics and Automation (ICRA)}, 95 | pages={8238--8244}, 96 | year={2022}, 97 | organization={IEEE} 98 | } 99 | ``` 100 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/group_points_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | 9 | #include "cuda_utils.h" 10 | 11 | // input: points(b, c, n) idx(b, npoints, nsample) 12 | // output: out(b, c, npoints, nsample) 13 | __global__ void group_points_kernel(int b, int c, int n, int npoints, 14 | int nsample, 15 | const float *__restrict__ points, 16 | const int *__restrict__ idx, 17 | float *__restrict__ out) { 18 | int batch_index = blockIdx.x; 19 | points += batch_index * n * c; 20 | idx += batch_index * npoints * nsample; 21 | out += batch_index * npoints * nsample * c; 22 | 23 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 24 | const int stride = blockDim.y * blockDim.x; 25 | for (int i = index; i < c * npoints; i += stride) { 26 | const int l = i / npoints; 27 | const int j = i % npoints; 28 | for (int k = 0; k < nsample; ++k) { 29 | int ii = idx[j * nsample + k]; 30 | out[(l * npoints + j) * nsample + k] = points[l * n + ii]; 31 | } 32 | } 33 | } 34 | 35 | void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample, 36 | const float *points, const int *idx, 37 | float *out) { 38 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 39 | 40 | group_points_kernel<<>>( 41 | b, c, n, npoints, nsample, points, idx, out); 42 | 43 | CUDA_CHECK_ERRORS(); 44 | } 45 | 46 | // input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample) 47 | // output: grad_points(b, c, n) 48 | __global__ void group_points_grad_kernel(int b, int c, int n, int npoints, 49 | int nsample, 50 | const float *__restrict__ grad_out, 51 | const int *__restrict__ idx, 52 | float *__restrict__ grad_points) { 53 | int batch_index = blockIdx.x; 54 | grad_out += batch_index * npoints * nsample * c; 55 | idx += batch_index * npoints * nsample; 56 | grad_points += batch_index * n * c; 57 | 58 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 59 | const int stride = blockDim.y * blockDim.x; 60 | for (int i = index; i < c * npoints; i += stride) { 61 | const int l = i / npoints; 62 | const int j = i % npoints; 63 | for (int k = 0; k < nsample; ++k) { 64 | int ii = idx[j * nsample + k]; 65 | atomicAdd(grad_points + l * n + ii, 66 | grad_out[(l * npoints + j) * nsample + k]); 67 | } 68 | } 69 | } 70 | 71 | void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 72 | int nsample, const float *grad_out, 73 | const int *idx, float *grad_points) { 74 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 75 | 76 | group_points_grad_kernel<<>>( 77 | b, c, n, npoints, nsample, grad_out, idx, grad_points); 78 | 79 | CUDA_CHECK_ERRORS(); 80 | } 81 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/sampling.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "sampling.h" 7 | #include "utils.h" 8 | 9 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 10 | const float *points, const int *idx, 11 | float *out); 12 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 13 | const float *grad_out, const int *idx, 14 | float *grad_points); 15 | 16 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 17 | const float *dataset, float *temp, 18 | int *idxs); 19 | 20 | at::Tensor gather_points(at::Tensor points, at::Tensor idx) { 21 | CHECK_CONTIGUOUS(points); 22 | CHECK_CONTIGUOUS(idx); 23 | CHECK_IS_FLOAT(points); 24 | CHECK_IS_INT(idx); 25 | 26 | if (points.type().is_cuda()) { 27 | CHECK_CUDA(idx); 28 | } 29 | 30 | at::Tensor output = 31 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 32 | at::device(points.device()).dtype(at::ScalarType::Float)); 33 | 34 | if (points.type().is_cuda()) { 35 | gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2), 36 | idx.size(1), points.data(), 37 | idx.data(), output.data()); 38 | } else { 39 | TORCH_CHECK(false, "CPU not supported"); 40 | } 41 | 42 | return output; 43 | } 44 | 45 | at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, 46 | const int n) { 47 | CHECK_CONTIGUOUS(grad_out); 48 | CHECK_CONTIGUOUS(idx); 49 | CHECK_IS_FLOAT(grad_out); 50 | CHECK_IS_INT(idx); 51 | 52 | if (grad_out.type().is_cuda()) { 53 | CHECK_CUDA(idx); 54 | } 55 | 56 | at::Tensor output = 57 | torch::zeros({grad_out.size(0), grad_out.size(1), n}, 58 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 59 | 60 | if (grad_out.type().is_cuda()) { 61 | gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n, 62 | idx.size(1), grad_out.data(), 63 | idx.data(), output.data()); 64 | } else { 65 | TORCH_CHECK(false, "CPU not supported"); 66 | } 67 | 68 | return output; 69 | } 70 | at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) { 71 | CHECK_CONTIGUOUS(points); 72 | CHECK_IS_FLOAT(points); 73 | 74 | at::Tensor output = 75 | torch::zeros({points.size(0), nsamples}, 76 | at::device(points.device()).dtype(at::ScalarType::Int)); 77 | 78 | at::Tensor tmp = 79 | torch::full({points.size(0), points.size(1)}, 1e10, 80 | at::device(points.device()).dtype(at::ScalarType::Float)); 81 | 82 | if (points.type().is_cuda()) { 83 | furthest_point_sampling_kernel_wrapper( 84 | points.size(0), points.size(1), nsamples, points.data(), 85 | tmp.data(), output.data()); 86 | } else { 87 | TORCH_CHECK(false, "CPU not supported"); 88 | } 89 | 90 | return output; 91 | } 92 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/interpolate.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include "interpolate.h" 7 | #include "utils.h" 8 | 9 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 10 | const float *known, float *dist2, int *idx); 11 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 12 | const float *points, const int *idx, 13 | const float *weight, float *out); 14 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 15 | const float *grad_out, 16 | const int *idx, const float *weight, 17 | float *grad_points); 18 | 19 | std::vector three_nn(at::Tensor unknowns, at::Tensor knows) { 20 | CHECK_CONTIGUOUS(unknowns); 21 | CHECK_CONTIGUOUS(knows); 22 | CHECK_IS_FLOAT(unknowns); 23 | CHECK_IS_FLOAT(knows); 24 | 25 | if (unknowns.type().is_cuda()) { 26 | CHECK_CUDA(knows); 27 | } 28 | 29 | at::Tensor idx = 30 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 31 | at::device(unknowns.device()).dtype(at::ScalarType::Int)); 32 | at::Tensor dist2 = 33 | torch::zeros({unknowns.size(0), unknowns.size(1), 3}, 34 | at::device(unknowns.device()).dtype(at::ScalarType::Float)); 35 | 36 | if (unknowns.type().is_cuda()) { 37 | three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1), 38 | unknowns.data(), knows.data(), 39 | dist2.data(), idx.data()); 40 | } else { 41 | TORCH_CHECK(false, "CPU not supported"); 42 | } 43 | 44 | return {dist2, idx}; 45 | } 46 | 47 | at::Tensor three_interpolate(at::Tensor points, at::Tensor idx, 48 | at::Tensor weight) { 49 | CHECK_CONTIGUOUS(points); 50 | CHECK_CONTIGUOUS(idx); 51 | CHECK_CONTIGUOUS(weight); 52 | CHECK_IS_FLOAT(points); 53 | CHECK_IS_INT(idx); 54 | CHECK_IS_FLOAT(weight); 55 | 56 | if (points.type().is_cuda()) { 57 | CHECK_CUDA(idx); 58 | CHECK_CUDA(weight); 59 | } 60 | 61 | at::Tensor output = 62 | torch::zeros({points.size(0), points.size(1), idx.size(1)}, 63 | at::device(points.device()).dtype(at::ScalarType::Float)); 64 | 65 | if (points.type().is_cuda()) { 66 | three_interpolate_kernel_wrapper( 67 | points.size(0), points.size(1), points.size(2), idx.size(1), 68 | points.data(), idx.data(), weight.data(), 69 | output.data()); 70 | } else { 71 | TORCH_CHECK(false, "CPU not supported"); 72 | } 73 | 74 | return output; 75 | } 76 | at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx, 77 | at::Tensor weight, const int m) { 78 | CHECK_CONTIGUOUS(grad_out); 79 | CHECK_CONTIGUOUS(idx); 80 | CHECK_CONTIGUOUS(weight); 81 | CHECK_IS_FLOAT(grad_out); 82 | CHECK_IS_INT(idx); 83 | CHECK_IS_FLOAT(weight); 84 | 85 | if (grad_out.type().is_cuda()) { 86 | CHECK_CUDA(idx); 87 | CHECK_CUDA(weight); 88 | } 89 | 90 | at::Tensor output = 91 | torch::zeros({grad_out.size(0), grad_out.size(1), m}, 92 | at::device(grad_out.device()).dtype(at::ScalarType::Float)); 93 | 94 | if (grad_out.type().is_cuda()) { 95 | three_interpolate_grad_kernel_wrapper( 96 | grad_out.size(0), grad_out.size(1), grad_out.size(2), m, 97 | grad_out.data(), idx.data(), weight.data(), 98 | output.data()); 99 | } else { 100 | TORCH_CHECK(false, "CPU not supported"); 101 | } 102 | 103 | return output; 104 | } 105 | -------------------------------------------------------------------------------- /models/decode.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import numpy as np 5 | import sys 6 | import os 7 | import time 8 | 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | ROOT_DIR = os.path.dirname(BASE_DIR) 11 | sys.path.append(ROOT_DIR) 12 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 13 | 14 | from utils.loss_utils import GRASP_MAX_WIDTH, THRESH_GOOD, THRESH_BAD,\ 15 | transform_point_cloud, generate_grasp_views,\ 16 | batch_viewpoint_params_to_matrix, huber_loss 17 | 18 | 19 | def index_select(index, input): 20 | N = input.shape[0] 21 | output = [] 22 | for i in range(N): 23 | x = index[i] 24 | out = input[i][x] 25 | output.append(out) 26 | output = torch.stack(output) 27 | return output 28 | 29 | 30 | def pred_decode(end_points): 31 | batch_size = len(end_points['point_clouds']) 32 | grasp_preds = [] 33 | for i in range(batch_size): 34 | ## load predictions 35 | objectness_score = end_points['objectness_score'][i].float() 36 | grasp_score = end_points['grasp_score_pred'][i].float() # 48, N 37 | grasp_score = grasp_score.transpose(0, 1).view(-1, 12, 4) # N, 12, 4 38 | grasp_center = end_points['fp2_xyz'][i].float() 39 | approaching = -end_points['grasp_top_view_xyz'][i].float() 40 | grasp_angle_class_score = end_points['grasp_angle_cls_pred'][i] # 12, N 41 | grasp_depth_class_score = end_points['grasp_depth_cls_pred'][i] # 4, N 42 | grasp_width = 1.2 * end_points['grasp_width_pred'][i] # 4, N 43 | grasp_width = torch.clamp(grasp_width, min=0, max=GRASP_MAX_WIDTH) 44 | 45 | # grasp angle 46 | grasp_angle_class = torch.argmax(grasp_angle_class_score, 0) # N 47 | grasp_angle = grasp_angle_class.float() / 12 * np.pi 48 | # grasp score & width & tolerance 49 | # grasp_angle_class_ = grasp_angle_class.unsqueeze(0) # 1, N 50 | 51 | grasp_score = index_select(grasp_angle_class, grasp_score) # N, 4 52 | # grasp_width = torch.gather(grasp_width, 0, grasp_angle_class_).squeeze(0) 53 | 54 | # grasp depth 55 | grasp_depth_class = torch.argmax(grasp_depth_class_score, 0) # N 56 | grasp_depth = (grasp_depth_class.float() + 1) * 0.01 57 | # grasp score & angle & width & tolerance 58 | grasp_score = index_select(grasp_depth_class, grasp_score) # N 59 | # grasp_angle = torch.gather(grasp_angle, 1, grasp_depth_class) 60 | grasp_width = index_select(grasp_depth_class, grasp_width.transpose(0, 1)) # N 61 | 62 | ## slice preds by objectness 63 | objectness_pred = torch.argmax(objectness_score, 0) 64 | objectness_mask = (objectness_pred == 1) 65 | grasp_score = grasp_score[objectness_mask].unsqueeze(-1) 66 | grasp_width = grasp_width[objectness_mask].unsqueeze(-1) 67 | grasp_depth = grasp_depth[objectness_mask].unsqueeze(-1) 68 | approaching = approaching[objectness_mask] 69 | grasp_angle = grasp_angle[objectness_mask] 70 | 71 | # objectness_label = end_points['objectness_label'] 72 | # fp2_inds = end_points['fp2_inds'].long() 73 | # objectness_label = torch.gather(objectness_label, 1, fp2_inds).squeeze(0) 74 | # objectness_mask1 = (objectness_label==1) 75 | 76 | grasp_center = grasp_center[objectness_mask] 77 | 78 | # grasp_score = grasp_score * grasp_tolerance / GRASP_MAX_TOLERANCE 79 | 80 | ## convert to rotation matrix 81 | Ns = grasp_angle.size(0) 82 | approaching_ = approaching.view(Ns, 3) 83 | grasp_angle_ = grasp_angle.view(Ns) 84 | rotation_matrix = batch_viewpoint_params_to_matrix(approaching_, grasp_angle_) 85 | rotation_matrix = rotation_matrix.view(Ns, 9) 86 | 87 | # merge preds 88 | grasp_height = 0.02 * torch.ones_like(grasp_score) 89 | obj_ids = -1 * torch.ones_like(grasp_score) 90 | grasp_preds.append( 91 | torch.cat([grasp_score, grasp_width, grasp_height, grasp_depth, rotation_matrix, grasp_center, obj_ids], 92 | axis=-1)) 93 | 94 | return grasp_preds -------------------------------------------------------------------------------- /utils/loss_utils.py: -------------------------------------------------------------------------------- 1 | """ Tools for loss computation. 2 | Author: chenxi-wang 3 | """ 4 | 5 | import torch 6 | import numpy as np 7 | 8 | GRASP_MAX_WIDTH = 0.1 9 | GRASP_MAX_TOLERANCE = 0.05 10 | THRESH_GOOD = 0.7 11 | THRESH_BAD = 0.1 12 | 13 | def transform_point_cloud(cloud, transform, format='4x4'): 14 | """ Transform points to new coordinates with transformation matrix. 15 | 16 | Input: 17 | cloud: [torch.FloatTensor, (N,3)] 18 | points in original coordinates 19 | transform: [torch.FloatTensor, (3,3)/(3,4)/(4,4)] 20 | transformation matrix, could be rotation only or rotation+translation 21 | format: [string, '3x3'/'3x4'/'4x4'] 22 | the shape of transformation matrix 23 | '3x3' --> rotation matrix 24 | '3x4'/'4x4' --> rotation matrix + translation matrix 25 | 26 | Output: 27 | cloud_transformed: [torch.FloatTensor, (N,3)] 28 | points in new coordinates 29 | """ 30 | if not (format == '3x3' or format == '4x4' or format == '3x4'): 31 | raise ValueError('Unknown transformation format, only support \'3x3\' or \'4x4\' or \'3x4\'.') 32 | if format == '3x3': 33 | cloud_transformed = torch.matmul(transform, cloud.T).T 34 | elif format == '4x4' or format == '3x4': 35 | ones = cloud.new_ones(cloud.size(0), device=cloud.device).unsqueeze(-1) 36 | cloud_ = torch.cat([cloud, ones], dim=1) 37 | cloud_transformed = torch.matmul(transform, cloud_.T).T 38 | cloud_transformed = cloud_transformed[:, :3] 39 | return cloud_transformed 40 | 41 | def generate_grasp_views(N=300, phi=(np.sqrt(5)-1)/2, center=np.zeros(3), r=1): 42 | """ View sampling on a unit sphere using Fibonacci lattices. 43 | Ref: https://arxiv.org/abs/0912.4540 44 | 45 | Input: 46 | N: [int] 47 | number of sampled views 48 | phi: [float] 49 | constant for view coordinate calculation, different phi's bring different distributions, default: (sqrt(5)-1)/2 50 | center: [np.ndarray, (3,), np.float32] 51 | sphere center 52 | r: [float] 53 | sphere radius 54 | 55 | Output: 56 | views: [torch.FloatTensor, (N,3)] 57 | sampled view coordinates 58 | """ 59 | views = [] 60 | for i in range(N): 61 | zi = (2 * i + 1) / N - 1 62 | xi = np.sqrt(1 - zi**2) * np.cos(2 * i * np.pi * phi) 63 | yi = np.sqrt(1 - zi**2) * np.sin(2 * i * np.pi * phi) 64 | views.append([xi, yi, zi]) 65 | views = r * np.array(views) + center 66 | return torch.from_numpy(views.astype(np.float32)) 67 | 68 | def batch_viewpoint_params_to_matrix(batch_towards, batch_angle): 69 | """ Transform approach vectors and in-plane rotation angles to rotation matrices. 70 | 71 | Input: 72 | batch_towards: [torch.FloatTensor, (N,3)] 73 | approach vectors in batch 74 | batch_angle: [torch.floatTensor, (N,)] 75 | in-plane rotation angles in batch 76 | 77 | Output: 78 | batch_matrix: [torch.floatTensor, (N,3,3)] 79 | rotation matrices in batch 80 | """ 81 | axis_x = batch_towards 82 | ones = torch.ones(axis_x.shape[0], dtype=axis_x.dtype, device=axis_x.device) 83 | zeros = torch.zeros(axis_x.shape[0], dtype=axis_x.dtype, device=axis_x.device) 84 | axis_y = torch.stack([-axis_x[:,1], axis_x[:,0], zeros], dim=-1) 85 | mask_y = (torch.norm(axis_y, dim=-1) == 0) 86 | axis_y[mask_y,1] = 1 87 | axis_x = axis_x / torch.norm(axis_x, dim=-1, keepdim=True) 88 | axis_y = axis_y / torch.norm(axis_y, dim=-1, keepdim=True) 89 | axis_z = torch.cross(axis_x, axis_y) 90 | sin = torch.sin(batch_angle) 91 | cos = torch.cos(batch_angle) 92 | R1 = torch.stack([ones, zeros, zeros, zeros, cos, -sin, zeros, sin, cos], dim=-1) 93 | R1 = R1.reshape([-1,3,3]) 94 | R2 = torch.stack([axis_x, axis_y, axis_z], dim=-1) 95 | batch_matrix = torch.matmul(R2, R1) 96 | return batch_matrix 97 | 98 | def huber_loss(error, delta=1.0): 99 | """ 100 | Args: 101 | error: Torch tensor (d1,d2,...,dk) 102 | Returns: 103 | loss: Torch tensor (d1,d2,...,dk) 104 | 105 | x = error = pred - gt or dist(pred,gt) 106 | 0.5 * |x|^2 if |x|<=d 107 | 0.5 * d^2 + d * (|x|-d) if |x|>d 108 | Author: Charles R. Qi 109 | Ref: https://github.com/charlesq34/frustum-pointnets/blob/master/models/model_util.py 110 | """ 111 | abs_error = torch.abs(error) 112 | quadratic = torch.clamp(abs_error, max=delta) 113 | linear = (abs_error - quadratic) 114 | loss = 0.5 * quadratic**2 + delta * linear 115 | return loss -------------------------------------------------------------------------------- /FGC_generate/rnn_neighbor.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d 2 | import numpy as np 3 | from sklearn import metrics 4 | from vector import orient_normals 5 | from tqdm import tqdm 6 | import os 7 | from contact_score import Contact_decision 8 | 9 | 10 | def load_grasp_labels(root, i): 11 | label = np.load(os.path.join(root, 'grasp_label', '{}_labels.npz'.format(str(i).zfill(3)))) 12 | obj_point = label['points'].astype(np.float32) 13 | 14 | return obj_point 15 | 16 | 17 | def get_vec_o3d(pcd): 18 | 19 | ''' 20 | input: PointCloud of object 21 | pcd.points: np.shape = (N, 3) 22 | :return: normal of object points 23 | outvec: np.shape = (N, 3) 24 | ''' 25 | 26 | pcd.estimate_normals() 27 | outvec = orient_normals(pcd) # return numpy with pc shape 28 | return outvec 29 | 30 | 31 | def search_radius_vector_3d(pcd, k, j, dis): 32 | ''' 33 | 34 | :param k: the number of search neighbor point 35 | :param j: the point index of object pc to search neighbor 36 | :param dis: the radius of search distance 37 | :return: give color to these points 38 | ''' 39 | pcd_tree = o3d.geometry.KDTreeFlann(pcd) 40 | # search_radius_vector_3d 半径近邻搜索 41 | pcd.colors[j] = [1, 0, 0] # 给定查询点并渲染为红色 42 | #[k1, idx1, _] = pcd_tree.search_radius_vector_3d(pcd.points[j], dis) 43 | [k1, idx1, _] = pcd_tree.search_knn_vector_3d(pcd.points[j], k) 44 | np.asarray(pcd.colors)[idx1[1:], :] = [0, 0, 1] 45 | 46 | 47 | def get_neigh_score(pcd): 48 | ''' 49 | 50 | :param vec: the normal vector of object pc 51 | np.shape = (N, 3) 52 | :return: point_curvity: np.shape = (N, 1) 53 | 54 | ''' 55 | vec = get_vec_o3d(pcd) 56 | pcd_tree = o3d.geometry.KDTreeFlann(pcd) 57 | k = 30 58 | dis = 0.02 59 | point_curvity = [] 60 | for j in range(len(pcd.points)): 61 | [k1, indx, _] = pcd_tree.search_knn_vector_3d(pcd.points[j], k) 62 | neigh_vec = vec[indx] 63 | #mse_neigh = mse(neigh_vec) 64 | cos_neigh = cosine(neigh_vec) 65 | if cos_neigh < 0: 66 | cos_neigh = 0.2 67 | point_curvity.append(cos_neigh) 68 | return np.asarray(point_curvity), vec 69 | 70 | 71 | def mse(neigh_vec): 72 | means = np.mean(neigh_vec, axis=0) 73 | mse = metrics.mean_squared_error(neigh_vec, np.tile(means, (neigh_vec.shape[0], 1))) 74 | return mse 75 | 76 | 77 | def cosine(neigh_vec): 78 | ''' 79 | 80 | :param neigh_vec: np.shape = (k, 3) 81 | neigh_vec[0] is the calculated point 82 | :return:cos_mean: float∈[-1, 1] 83 | the mean cosine distance of raw point in the neighbor 84 | ''' 85 | raw = neigh_vec[0] 86 | raw = np.expand_dims(raw, axis=1) 87 | others = neigh_vec[1:] 88 | cos = np.dot(others, raw) # 内积距离 89 | cos_mean = np.mean(cos) # 邻域内各点到查询点的平均内积距离 90 | return cos_mean 91 | 92 | 93 | def vis_search(pcd): 94 | vec = get_vec_o3d(pcd) 95 | pcd.normals = o3d.utility.Vector3dVector(vec) 96 | o3d.visualization.draw_geometries([pcd], width=800, height=600, left=50, top=50, point_show_normal=True) 97 | 98 | pcd.paint_uniform_color([0.5, 0.5, 0.5]) 99 | point_curvity = get_neigh_score(pcd) 100 | arg = np.argmax(point_curvity) 101 | 102 | search_radius_vector_3d(pcd, 20, arg, 0.01) 103 | o3d.visualization.draw_geometries([pcd], window_name='the max cosine', width=800, height=600, left=50, top=50, 104 | point_show_normal=False, mesh_show_wireframe=False, 105 | mesh_show_back_face=False) 106 | 107 | 108 | if __name__=='__main__': 109 | root = '../grasp_data' 110 | # for i in tqdm(range(88)): 111 | # obj_pc = load_grasp_labels(root, i) 112 | # pcd = o3d.geometry.PointCloud() 113 | # pcd.points = o3d.utility.Vector3dVector(obj_pc) 114 | # score3, normals = get_neigh_score(pcd) # score3:(N, ) normals:(N, 3) 115 | # savepath = os.path.join(root, 'normals_score', '{}_labels.npz'.format(str(i).zfill(3))) 116 | # np.savez(savepath, normals=normals, score3=score3) 117 | # print(score3) 118 | 119 | eg = np.load(os.path.join(root, 'normals_score', '{}_labels.npz'.format(str(50).zfill(3)))) 120 | print(eg.files) 121 | normals = eg['normals'] 122 | score3 = eg['score3'] 123 | print('1') 124 | 125 | #contact_decision = Contact_decision(root) 126 | #pc, _, four, _ = contact_decision.contact_find() 127 | #left, right, _, _ = four 128 | #curvity = Curvity_score(pc, left, right) 129 | #curvity.vis_search() 130 | #score3 = curvity.get_neigh_score() 131 | # for i in range(pc.shape[0]): 132 | # if np.any(pc[i]==left): 133 | # left_ind = i 134 | # elif np.any(pc[i]==right): 135 | # right_ind = i 136 | # score_of_contact = (score3[left_ind]+score3[right_ind])/2 137 | # print(score_of_contact) 138 | 139 | 140 | 141 | 142 | 143 | 144 | -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/interpolate_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | #include 9 | 10 | #include "cuda_utils.h" 11 | 12 | // input: unknown(b, n, 3) known(b, m, 3) 13 | // output: dist2(b, n, 3), idx(b, n, 3) 14 | __global__ void three_nn_kernel(int b, int n, int m, 15 | const float *__restrict__ unknown, 16 | const float *__restrict__ known, 17 | float *__restrict__ dist2, 18 | int *__restrict__ idx) { 19 | int batch_index = blockIdx.x; 20 | unknown += batch_index * n * 3; 21 | known += batch_index * m * 3; 22 | dist2 += batch_index * n * 3; 23 | idx += batch_index * n * 3; 24 | 25 | int index = threadIdx.x; 26 | int stride = blockDim.x; 27 | for (int j = index; j < n; j += stride) { 28 | float ux = unknown[j * 3 + 0]; 29 | float uy = unknown[j * 3 + 1]; 30 | float uz = unknown[j * 3 + 2]; 31 | 32 | double best1 = 1e40, best2 = 1e40, best3 = 1e40; 33 | int besti1 = 0, besti2 = 0, besti3 = 0; 34 | for (int k = 0; k < m; ++k) { 35 | float x = known[k * 3 + 0]; 36 | float y = known[k * 3 + 1]; 37 | float z = known[k * 3 + 2]; 38 | float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z); 39 | if (d < best1) { 40 | best3 = best2; 41 | besti3 = besti2; 42 | best2 = best1; 43 | besti2 = besti1; 44 | best1 = d; 45 | besti1 = k; 46 | } else if (d < best2) { 47 | best3 = best2; 48 | besti3 = besti2; 49 | best2 = d; 50 | besti2 = k; 51 | } else if (d < best3) { 52 | best3 = d; 53 | besti3 = k; 54 | } 55 | } 56 | dist2[j * 3 + 0] = best1; 57 | dist2[j * 3 + 1] = best2; 58 | dist2[j * 3 + 2] = best3; 59 | 60 | idx[j * 3 + 0] = besti1; 61 | idx[j * 3 + 1] = besti2; 62 | idx[j * 3 + 2] = besti3; 63 | } 64 | } 65 | 66 | void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown, 67 | const float *known, float *dist2, int *idx) { 68 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 69 | three_nn_kernel<<>>(b, n, m, unknown, known, 70 | dist2, idx); 71 | 72 | CUDA_CHECK_ERRORS(); 73 | } 74 | 75 | // input: points(b, c, m), idx(b, n, 3), weight(b, n, 3) 76 | // output: out(b, c, n) 77 | __global__ void three_interpolate_kernel(int b, int c, int m, int n, 78 | const float *__restrict__ points, 79 | const int *__restrict__ idx, 80 | const float *__restrict__ weight, 81 | float *__restrict__ out) { 82 | int batch_index = blockIdx.x; 83 | points += batch_index * m * c; 84 | 85 | idx += batch_index * n * 3; 86 | weight += batch_index * n * 3; 87 | 88 | out += batch_index * n * c; 89 | 90 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 91 | const int stride = blockDim.y * blockDim.x; 92 | for (int i = index; i < c * n; i += stride) { 93 | const int l = i / n; 94 | const int j = i % n; 95 | float w1 = weight[j * 3 + 0]; 96 | float w2 = weight[j * 3 + 1]; 97 | float w3 = weight[j * 3 + 2]; 98 | 99 | int i1 = idx[j * 3 + 0]; 100 | int i2 = idx[j * 3 + 1]; 101 | int i3 = idx[j * 3 + 2]; 102 | 103 | out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 + 104 | points[l * m + i3] * w3; 105 | } 106 | } 107 | 108 | void three_interpolate_kernel_wrapper(int b, int c, int m, int n, 109 | const float *points, const int *idx, 110 | const float *weight, float *out) { 111 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 112 | three_interpolate_kernel<<>>( 113 | b, c, m, n, points, idx, weight, out); 114 | 115 | CUDA_CHECK_ERRORS(); 116 | } 117 | 118 | // input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3) 119 | // output: grad_points(b, c, m) 120 | 121 | __global__ void three_interpolate_grad_kernel( 122 | int b, int c, int n, int m, const float *__restrict__ grad_out, 123 | const int *__restrict__ idx, const float *__restrict__ weight, 124 | float *__restrict__ grad_points) { 125 | int batch_index = blockIdx.x; 126 | grad_out += batch_index * n * c; 127 | idx += batch_index * n * 3; 128 | weight += batch_index * n * 3; 129 | grad_points += batch_index * m * c; 130 | 131 | const int index = threadIdx.y * blockDim.x + threadIdx.x; 132 | const int stride = blockDim.y * blockDim.x; 133 | for (int i = index; i < c * n; i += stride) { 134 | const int l = i / n; 135 | const int j = i % n; 136 | float w1 = weight[j * 3 + 0]; 137 | float w2 = weight[j * 3 + 1]; 138 | float w3 = weight[j * 3 + 2]; 139 | 140 | int i1 = idx[j * 3 + 0]; 141 | int i2 = idx[j * 3 + 1]; 142 | int i3 = idx[j * 3 + 2]; 143 | 144 | atomicAdd(grad_points + l * m + i1, grad_out[i] * w1); 145 | atomicAdd(grad_points + l * m + i2, grad_out[i] * w2); 146 | atomicAdd(grad_points + l * m + i3, grad_out[i] * w3); 147 | } 148 | } 149 | 150 | void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m, 151 | const float *grad_out, 152 | const int *idx, const float *weight, 153 | float *grad_points) { 154 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 155 | three_interpolate_grad_kernel<<>>( 156 | b, c, n, m, grad_out, idx, weight, grad_points); 157 | 158 | CUDA_CHECK_ERRORS(); 159 | } 160 | -------------------------------------------------------------------------------- /demo.py: -------------------------------------------------------------------------------- 1 | """ Demo to show prediction results. 2 | Author: chenxi-wang 3 | """ 4 | 5 | import os 6 | import sys 7 | import numpy as np 8 | import open3d as o3d 9 | import argparse 10 | import importlib 11 | import scipy.io as scio 12 | from PIL import Image 13 | import cv2 14 | import torch 15 | from graspnetAPI import GraspGroup 16 | 17 | ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) 18 | sys.path.append(os.path.join(ROOT_DIR, 'models')) 19 | sys.path.append(os.path.join(ROOT_DIR, 'dataset')) 20 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 21 | 22 | from models.FGC_graspnet import FGC_graspnet 23 | from dataset.graspnet_dataset import GraspNetDataset 24 | from utils.collision_detector import ModelFreeCollisionDetector 25 | from utils.data_utils import CameraInfo, create_point_cloud_from_depth_image 26 | from models.decode import pred_decode 27 | 28 | parser = argparse.ArgumentParser() 29 | parser.add_argument('--checkpoint_path', default='./logs/baseline/checkpoint.tar', help='Model checkpoint path') 30 | parser.add_argument('--num_point', type=int, default=20000, help='Point Number [default: 20000]') 31 | parser.add_argument('--num_view', type=int, default=300, help='View Number [default: 300]') 32 | parser.add_argument('--collision_thresh', type=float, default=0.01, help='Collision Threshold in collision detection [default: 0.01]') 33 | parser.add_argument('--voxel_size', type=float, default=0.01, help='Voxel Size to process point clouds before collision detection [default: 0.01]') 34 | cfgs = parser.parse_args() 35 | 36 | 37 | def get_net(): 38 | # Init the model 39 | net = FGC_graspnet(input_feature_dim=0, num_view=cfgs.num_view, num_angle=12, num_depth=4, 40 | cylinder_radius=0.05, hmin=-0.02, hmax=0.02, is_training=False, is_demo=True) 41 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 42 | net.to(device) 43 | # Load checkpoint 44 | checkpoint = torch.load(cfgs.checkpoint_path) 45 | net.load_state_dict(checkpoint['model_state_dict']) 46 | start_epoch = checkpoint['epoch'] 47 | print("-> loaded checkpoint %s (epoch: %d)"%(cfgs.checkpoint_path, start_epoch)) 48 | # set model to eval mode 49 | net.eval() 50 | return net 51 | 52 | 53 | def get_and_process_data(data_dir): 54 | # load data 55 | color = np.array(Image.open(os.path.join(data_dir, 'orange_color.png')).resize((320, 240)), dtype=np.float32) / 255.0 56 | raw = np.fromfile(os.path.join(data_dir, 'orange_depth.raw'), dtype='uint16') 57 | depth = np.resize(raw, (240, 320)) 58 | workspace_mask = np.array(Image.open(os.path.join(data_dir, 'workspace_mask.png')).resize((320, 240))) 59 | #meta = scio.loadmat(os.path.join(data_dir, 'meta.mat')) 60 | #intrinsic = meta['intrinsic_matrix'] 61 | #factor_depth = meta['factor_depth'] 62 | 63 | # generate cloud 64 | '''we use the intrinsic of the Realsense D435i camera in our experiments, 65 | you can change the intrinsic by yourself. 66 | ''' 67 | camera = CameraInfo(320.0, 240.0, 296.282990, 296.282990, 159.566162, 123.468735, 1000) 68 | cloud = create_point_cloud_from_depth_image(depth, camera, organized=True) 69 | 70 | # get valid points 71 | mask = (workspace_mask & (depth > 0)) 72 | cloud_masked = cloud[mask] 73 | color_masked = color[mask] 74 | 75 | # sample points 76 | if len(cloud_masked) >= cfgs.num_point: 77 | idxs = np.random.choice(len(cloud_masked), cfgs.num_point, replace=False) 78 | else: 79 | idxs1 = np.arange(len(cloud_masked)) 80 | idxs2 = np.random.choice(len(cloud_masked), cfgs.num_point-len(cloud_masked), replace=True) 81 | idxs = np.concatenate([idxs1, idxs2], axis=0) 82 | 83 | cloud_sampled = cloud_masked[idxs] 84 | color_sampled = color_masked[idxs] 85 | #cloud_sampled = xyzrgb[:, 0:3] 86 | #color_sampled = xyzrgb[:, 3:6] 87 | 88 | # convert data 89 | cloud = o3d.geometry.PointCloud() 90 | cloud.points = o3d.utility.Vector3dVector(cloud_sampled.astype(np.float32)) 91 | cloud.colors = o3d.utility.Vector3dVector(color_sampled.astype(np.float32)) 92 | end_points = dict() 93 | cloud_sampled = torch.from_numpy(cloud_sampled[np.newaxis].astype(np.float32)) 94 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 95 | cloud_sampled = cloud_sampled.to(device) 96 | end_points['point_clouds'] = cloud_sampled 97 | end_points['cloud_colors'] = color_sampled 98 | 99 | return end_points, cloud 100 | 101 | 102 | def get_grasps(net, end_points): 103 | # Forward pass 104 | with torch.no_grad(): 105 | end_points = net(end_points) 106 | grasp_preds = pred_decode(end_points) 107 | gg_array = grasp_preds[0].detach().cpu().numpy() 108 | gg = GraspGroup(gg_array) 109 | # if save the results, return gg_array 110 | return gg 111 | 112 | 113 | def collision_detection(gg, cloud): 114 | mfcdetector = ModelFreeCollisionDetector(cloud, voxel_size=cfgs.voxel_size) 115 | collision_mask = mfcdetector.detect(gg, approach_dist=0.05, collision_thresh=cfgs.collision_thresh) 116 | gg = gg[~collision_mask] 117 | return gg 118 | 119 | 120 | def vis_grasps_save(gg_array, cloud): 121 | np.save('doc/grip.npy', gg_array) 122 | o3d.io.write_point_cloud('doc/cloud.ply', cloud) 123 | 124 | 125 | def vis_grasps(gg, cloud): 126 | gg.nms() 127 | gg.sort_by_score() 128 | gg = gg[:50] 129 | grippers = gg.to_open3d_geometry_list() 130 | o3d.visualization.draw_geometries([cloud, *grippers]) 131 | 132 | 133 | def demo(data_dir): 134 | net = get_net() 135 | end_points, cloud = get_and_process_data(data_dir) 136 | gg = get_grasps(net, end_points) 137 | #if cfgs.collision_thresh > 0: 138 | # gg = collision_detection(gg, np.array(cloud.points)) 139 | vis_grasps(gg, cloud) 140 | 141 | 142 | if __name__=='__main__': 143 | data_dir = 'doc/example_data' 144 | demo(data_dir) 145 | -------------------------------------------------------------------------------- /utils/data_utils.py: -------------------------------------------------------------------------------- 1 | """ Tools for data processing. 2 | Author: chenxi-wang 3 | """ 4 | 5 | import numpy as np 6 | 7 | class CameraInfo(): 8 | """ Camera intrisics for point cloud creation. """ 9 | def __init__(self, width, height, fx, fy, cx, cy, scale): 10 | self.width = width 11 | self.height = height 12 | self.fx = fx 13 | self.fy = fy 14 | self.cx = cx 15 | self.cy = cy 16 | self.scale = scale 17 | 18 | def create_point_cloud_from_depth_image(depth, camera, organized=True): 19 | """ Generate point cloud using depth image only. 20 | 21 | Input: 22 | depth: [numpy.ndarray, (H,W), numpy.float32] 23 | depth image 24 | camera: [CameraInfo] 25 | camera intrinsics 26 | organized: bool 27 | whether to keep the cloud in image shape (H,W,3) 28 | 29 | Output: 30 | cloud: [numpy.ndarray, (H,W,3)/(H*W,3), numpy.float32] 31 | generated cloud, (H,W,3) for organized=True, (H*W,3) for organized=False 32 | """ 33 | assert(depth.shape[0] == camera.height and depth.shape[1] == camera.width) 34 | xmap = np.arange(camera.width) 35 | ymap = np.arange(camera.height) 36 | xmap, ymap = np.meshgrid(xmap, ymap) 37 | points_z = depth / camera.scale 38 | points_x = (xmap - camera.cx) * points_z / camera.fx 39 | points_y = (ymap - camera.cy) * points_z / camera.fy 40 | cloud = np.stack([points_x, points_y, points_z], axis=-1) 41 | if not organized: 42 | cloud = cloud.reshape([-1, 3]) 43 | return cloud 44 | 45 | def transform_point_cloud(cloud, transform, format='4x4'): 46 | """ Transform points to new coordinates with transformation matrix. 47 | 48 | Input: 49 | cloud: [np.ndarray, (N,3), np.float32] 50 | points in original coordinates 51 | transform: [np.ndarray, (3,3)/(3,4)/(4,4), np.float32] 52 | transformation matrix, could be rotation only or rotation+translation 53 | format: [string, '3x3'/'3x4'/'4x4'] 54 | the shape of transformation matrix 55 | '3x3' --> rotation matrix 56 | '3x4'/'4x4' --> rotation matrix + translation matrix 57 | 58 | Output: 59 | cloud_transformed: [np.ndarray, (N,3), np.float32] 60 | points in new coordinates 61 | """ 62 | if not (format == '3x3' or format == '4x4' or format == '3x4'): 63 | raise ValueError('Unknown transformation format, only support \'3x3\' or \'4x4\' or \'3x4\'.') 64 | if format == '3x3': 65 | cloud_transformed = np.dot(transform, cloud.T).T 66 | elif format == '4x4' or format == '3x4': 67 | ones = np.ones(cloud.shape[0])[:, np.newaxis] 68 | cloud_ = np.concatenate([cloud, ones], axis=1) 69 | cloud_transformed = np.dot(transform, cloud_.T).T 70 | cloud_transformed = cloud_transformed[:, :3] 71 | return cloud_transformed 72 | 73 | def compute_point_dists(A, B): 74 | """ Compute pair-wise point distances in two matrices. 75 | 76 | Input: 77 | A: [np.ndarray, (N,3), np.float32] 78 | point cloud A 79 | B: [np.ndarray, (M,3), np.float32] 80 | point cloud B 81 | 82 | Output: 83 | dists: [np.ndarray, (N,M), np.float32] 84 | distance matrix 85 | """ 86 | A = A[:, np.newaxis, :] 87 | B = B[np.newaxis, :, :] 88 | dists = np.linalg.norm(A-B, axis=-1) 89 | return dists 90 | 91 | def remove_invisible_grasp_points(cloud, grasp_points, pose, th=0.01): 92 | """ Remove invisible part of object model according to scene point cloud. 93 | 94 | Input: 95 | cloud: [np.ndarray, (N,3), np.float32] 96 | scene point cloud 97 | grasp_points: [np.ndarray, (M,3), np.float32] 98 | grasp point label in object coordinates 99 | pose: [np.ndarray, (4,4), np.float32] 100 | transformation matrix from object coordinates to world coordinates 101 | th: [float] 102 | if the minimum distance between a grasp point and the scene points is greater than outlier, the point will be removed 103 | 104 | Output: 105 | visible_mask: [np.ndarray, (M,), np.bool] 106 | mask to show the visible part of grasp points 107 | """ 108 | grasp_points_trans = transform_point_cloud(grasp_points, pose) 109 | dists = compute_point_dists(grasp_points_trans, cloud) 110 | min_dists = dists.min(axis=1) 111 | visible_mask = (min_dists < th) 112 | return visible_mask 113 | 114 | def get_workspace_mask(cloud, seg, trans=None, organized=True, outlier=0): 115 | """ Keep points in workspace as input. 116 | 117 | Input: 118 | cloud: [np.ndarray, (H,W,3), np.float32] 119 | scene point cloud 120 | seg: [np.ndarray, (H,W,), np.uint8] 121 | segmantation label of scene points 122 | trans: [np.ndarray, (4,4), np.float32] 123 | transformation matrix for scene points, default: None. 124 | organized: [bool] 125 | whether to keep the cloud in image shape (H,W,3) 126 | outlier: [float] 127 | if the distance between a point and workspace is greater than outlier, the point will be removed 128 | 129 | Output: 130 | workspace_mask: [np.ndarray, (H,W)/(H*W,), np.bool] 131 | mask to indicate whether scene points are in workspace 132 | """ 133 | if organized: 134 | h, w, _ = cloud.shape 135 | cloud = cloud.reshape([h*w, 3]) 136 | seg = seg.reshape(h*w) 137 | if trans is not None: 138 | cloud = transform_point_cloud(cloud, trans) 139 | foreground = cloud[seg>0] 140 | xmin, ymin, zmin = foreground.min(axis=0) 141 | xmax, ymax, zmax = foreground.max(axis=0) 142 | mask_x = ((cloud[:,0] > xmin-outlier) & (cloud[:,0] < xmax+outlier)) 143 | mask_y = ((cloud[:,1] > ymin-outlier) & (cloud[:,1] < ymax+outlier)) 144 | mask_z = ((cloud[:,2] > zmin-outlier) & (cloud[:,2] < zmax+outlier)) 145 | workspace_mask = (mask_x & mask_y & mask_z) 146 | if organized: 147 | workspace_mask = workspace_mask.reshape([h, w]) 148 | 149 | return workspace_mask -------------------------------------------------------------------------------- /models/backbone.py: -------------------------------------------------------------------------------- 1 | """ PointNet2 backbone for feature learning. 2 | Author: Charles R. Qi 3 | """ 4 | import os 5 | import sys 6 | import torch 7 | import torch.nn as nn 8 | from knn.knn_modules import knn 9 | 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | ROOT_DIR = os.path.dirname(BASE_DIR) 12 | sys.path.append(ROOT_DIR) 13 | sys.path.append(os.path.join(ROOT_DIR, 'pointnet2')) 14 | 15 | from pointnet2.pointnet2_modules import PointnetSAModuleVotes, PointnetFPModule 16 | 17 | 18 | class Pointnet2Backbone(nn.Module): 19 | r""" 20 | Backbone network for point cloud feature learning. 21 | Based on Pointnet++ single-scale grouping network. 22 | 23 | Parameters 24 | ---------- 25 | input_feature_dim: int 26 | Number of input channels in the feature descriptor for each point. 27 | e.g. 3 for RGB. 28 | """ 29 | 30 | def __init__(self, input_feature_dim=0): 31 | super().__init__() 32 | 33 | self.sa1 = PointnetSAModuleVotes( 34 | npoint=2048, 35 | radius=0.04, 36 | nsample=64, 37 | mlp=[input_feature_dim, 64, 64, 128], 38 | use_xyz=True, 39 | normalize_xyz=True 40 | ) 41 | 42 | self.sa2 = PointnetSAModuleVotes( 43 | npoint=1024, 44 | radius=0.1, 45 | nsample=32, 46 | mlp=[128, 128, 128, 256], 47 | use_xyz=True, 48 | normalize_xyz=True 49 | ) 50 | 51 | self.sa3 = PointnetSAModuleVotes( 52 | npoint=512, 53 | radius=0.2, 54 | nsample=16, 55 | mlp=[256, 128, 128, 256], 56 | use_xyz=True, 57 | normalize_xyz=True 58 | ) 59 | 60 | self.sa4 = PointnetSAModuleVotes( 61 | npoint=256, 62 | radius=0.3, 63 | nsample=16, 64 | mlp=[256, 128, 128, 256], 65 | use_xyz=True, 66 | normalize_xyz=True 67 | ) 68 | 69 | self.fp1 = PointnetFPModule(mlp=[256+256,256,256]) 70 | self.fp2 = PointnetFPModule(mlp=[256+256,256,256]) 71 | 72 | #self.global_att = StackedAttention() 73 | 74 | def _break_up_pc(self, pc): 75 | xyz = pc[..., 0:3].contiguous() 76 | features = ( 77 | pc[..., 3:].transpose(1, 2).contiguous() 78 | if pc.size(-1) > 3 else None 79 | ) 80 | 81 | return xyz, features 82 | 83 | def forward(self, pointcloud: torch.cuda.FloatTensor, end_points=None): 84 | r""" 85 | Forward pass of the network 86 | 87 | Parameters 88 | ---------- 89 | pointcloud: Variable(torch.cuda.FloatTensor) 90 | (B, N, 3 + input_feature_dim) tensor 91 | Point cloud to run predicts on 92 | Each point in the point-cloud MUST 93 | be formated as (x, y, z, features...) 94 | 95 | Returns 96 | ---------- 97 | end_points: {XXX_xyz, XXX_features, XXX_inds} 98 | XXX_xyz: float32 Tensor of shape (B,K,3) 99 | XXX_features: float32 Tensor of shape (B,D,K) 100 | XXX_inds: int64 Tensor of shape (B,K) values in [0,N-1] 101 | """ 102 | if not end_points: end_points = {} 103 | batch_size = pointcloud.shape[0] 104 | 105 | xyz, features = self._break_up_pc(pointcloud) 106 | end_points['input_xyz'] = xyz 107 | end_points['input_features'] = features 108 | 109 | # --------- 4 SET ABSTRACTION LAYERS --------- 110 | xyz, features, fps_inds = self.sa1(xyz, features) 111 | end_points['sa1_inds'] = fps_inds 112 | end_points['sa1_xyz'] = xyz 113 | end_points['sa1_features'] = features 114 | 115 | xyz, features, fps_inds = self.sa2(xyz, features) # this fps_inds is just 0,1,...,1023 116 | end_points['sa2_inds'] = fps_inds 117 | end_points['sa2_xyz'] = xyz 118 | end_points['sa2_features'] = features 119 | 120 | xyz, features, fps_inds = self.sa3(xyz, features) # this fps_inds is just 0,1,...,255 121 | end_points['sa3_xyz'] = xyz 122 | end_points['sa3_features'] = features 123 | 124 | xyz, features, fps_inds = self.sa4(xyz, features) # this fps_inds is just 0,1,...,255 125 | end_points['sa4_xyz'] = xyz 126 | end_points['sa4_features'] = features 127 | 128 | 129 | # --------- 2 FEATURE UPSAMPLING LAYERS -------- 130 | features = self.fp1(end_points['sa3_xyz'], end_points['sa4_xyz'], end_points['sa3_features'], 131 | end_points['sa4_features']) 132 | features = self.fp2(end_points['sa2_xyz'], end_points['sa3_xyz'], end_points['sa2_features'], features) 133 | #features = self.global_att(features) 134 | end_points['fp2_features'] = features 135 | 136 | end_points['fp2_xyz'] = end_points['sa2_xyz'] 137 | num_seed = end_points['fp2_xyz'].shape[1] 138 | end_points['fp2_inds'] = end_points['sa1_inds'][:, 0:num_seed] # indices among the entire input point clouds 139 | 140 | return features, end_points['fp2_xyz'], end_points 141 | 142 | 143 | class Local_attention(nn.Module): 144 | def __init__(self, channels): 145 | super().__init__() 146 | self.q_conv = nn.Conv1d(channels, channels // 4, 1, bias=False) 147 | self.k_conv = nn.Conv1d(channels, channels // 4, 1, bias=False) 148 | self.q_conv.weight = self.k_conv.weight 149 | self.v_conv = nn.Conv1d(channels, channels, 1) 150 | self.trans_conv = nn.Conv1d(channels, channels, 1) 151 | self.after_norm = nn.BatchNorm1d(channels) 152 | self.act = nn.ReLU() 153 | self.softmax = nn.Softmax(dim=-1) 154 | 155 | def forward(self, x): 156 | # x.shape b*n, c, k 157 | x_q = self.q_conv(x).permute(0, 2, 1) # b*n, k, c 158 | x_k = self.k_conv(x) # b*n, c, k 159 | x_v = self.v_conv(x) # b*n, c, k 160 | energy = x_q @ x_k # b*n, k, k 161 | attention = self.softmax(energy) 162 | attention = attention / (1e-9 + attention.sum(dim=2, keepdims=True)) 163 | x_r = x_v @ attention # b*n, c, k 164 | x_r = self.act(self.after_norm(self.trans_conv(x - x_r))) 165 | x = x + x_r 166 | return x 167 | 168 | 169 | 170 | 171 | 172 | 173 | -------------------------------------------------------------------------------- /utils/collision_detector.py: -------------------------------------------------------------------------------- 1 | """ Collision detection to remove collided grasp pose predictions. 2 | Author: chenxi-wang 3 | """ 4 | 5 | import os 6 | import sys 7 | import numpy as np 8 | import open3d as o3d 9 | 10 | class ModelFreeCollisionDetector(): 11 | """ Collision detection in scenes without object labels. Current finger width and length are fixed. 12 | 13 | Input: 14 | scene_points: [numpy.ndarray, (N,3), numpy.float32] 15 | the scene points to detect 16 | voxel_size: [float] 17 | used for downsample 18 | 19 | Example usage: 20 | mfcdetector = ModelFreeCollisionDetector(scene_points, voxel_size=0.005) 21 | collision_mask = mfcdetector.detect(grasp_group, approach_dist=0.03) 22 | collision_mask, iou_list = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05, return_ious=True) 23 | collision_mask, empty_mask = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05, 24 | return_empty_grasp=True, empty_thresh=0.01) 25 | collision_mask, empty_mask, iou_list = mfcdetector.detect(grasp_group, approach_dist=0.03, collision_thresh=0.05, 26 | return_empty_grasp=True, empty_thresh=0.01, return_ious=True) 27 | """ 28 | def __init__(self, scene_points, voxel_size=0.005): 29 | self.finger_width = 0.01 30 | self.finger_length = 0.06 31 | self.voxel_size = voxel_size 32 | scene_cloud = o3d.geometry.PointCloud() 33 | scene_cloud.points = o3d.utility.Vector3dVector(scene_points) 34 | scene_cloud = scene_cloud.voxel_down_sample(voxel_size) 35 | self.scene_points = np.array(scene_cloud.points) 36 | 37 | def detect(self, grasp_group, approach_dist=0.03, collision_thresh=0.05, return_empty_grasp=False, empty_thresh=0.01, return_ious=False): 38 | """ Detect collision of grasps. 39 | 40 | Input: 41 | grasp_group: [GraspGroup, M grasps] 42 | the grasps to check 43 | approach_dist: [float] 44 | the distance for a gripper to move along approaching direction before grasping 45 | this shifting space requires no point either 46 | collision_thresh: [float] 47 | if global collision iou is greater than this threshold, 48 | a collision is detected 49 | return_empty_grasp: [bool] 50 | if True, return a mask to imply whether there are objects in a grasp 51 | empty_thresh: [float] 52 | if inner space iou is smaller than this threshold, 53 | a collision is detected 54 | only set when [return_empty_grasp] is True 55 | return_ious: [bool] 56 | if True, return global collision iou and part collision ious 57 | 58 | Output: 59 | collision_mask: [numpy.ndarray, (M,), numpy.bool] 60 | True implies collision 61 | [optional] empty_mask: [numpy.ndarray, (M,), numpy.bool] 62 | True implies empty grasp 63 | only returned when [return_empty_grasp] is True 64 | [optional] iou_list: list of [numpy.ndarray, (M,), numpy.float32] 65 | global and part collision ious, containing 66 | [global_iou, left_iou, right_iou, bottom_iou, shifting_iou] 67 | only returned when [return_ious] is True 68 | """ 69 | approach_dist = max(approach_dist, self.finger_width) 70 | T = grasp_group.translations 71 | R = grasp_group.rotation_matrices 72 | heights = grasp_group.heights[:,np.newaxis] 73 | depths = grasp_group.depths[:,np.newaxis] 74 | widths = grasp_group.widths[:,np.newaxis] 75 | targets = self.scene_points[np.newaxis,:,:] - T[:,np.newaxis,:] 76 | targets = np.matmul(targets, R) 77 | 78 | ## collision detection 79 | # height mask 80 | mask1 = ((targets[:,:,2] > -heights/2) & (targets[:,:,2] < heights/2)) 81 | # left finger mask 82 | mask2 = ((targets[:,:,0] > depths - self.finger_length) & (targets[:,:,0] < depths)) 83 | mask3 = (targets[:,:,1] > -(widths/2 + self.finger_width)) 84 | mask4 = (targets[:,:,1] < -widths/2) 85 | # right finger mask 86 | mask5 = (targets[:,:,1] < (widths/2 + self.finger_width)) 87 | mask6 = (targets[:,:,1] > widths/2) 88 | # bottom mask 89 | mask7 = ((targets[:,:,0] <= depths - self.finger_length)\ 90 | & (targets[:,:,0] > depths - self.finger_length - self.finger_width)) 91 | # shifting mask 92 | mask8 = ((targets[:,:,0] <= depths - self.finger_length - self.finger_width)\ 93 | & (targets[:,:,0] > depths - self.finger_length - self.finger_width - approach_dist)) 94 | 95 | # get collision mask of each point 96 | left_mask = (mask1 & mask2 & mask3 & mask4) 97 | right_mask = (mask1 & mask2 & mask5 & mask6) 98 | bottom_mask = (mask1 & mask3 & mask5 & mask7) 99 | shifting_mask = (mask1 & mask3 & mask5 & mask8) 100 | global_mask = (left_mask | right_mask | bottom_mask | shifting_mask) 101 | 102 | # calculate equivalant volume of each part 103 | left_right_volume = (heights * self.finger_length * self.finger_width / (self.voxel_size**3)).reshape(-1) 104 | bottom_volume = (heights * (widths+2*self.finger_width) * self.finger_width / (self.voxel_size**3)).reshape(-1) 105 | shifting_volume = (heights * (widths+2*self.finger_width) * approach_dist / (self.voxel_size**3)).reshape(-1) 106 | volume = left_right_volume*2 + bottom_volume + shifting_volume 107 | 108 | # get collision iou of each part 109 | global_iou = global_mask.sum(axis=1) / (volume+1e-6) 110 | 111 | # get collison mask 112 | collision_mask = (global_iou > collision_thresh) 113 | 114 | if not (return_empty_grasp or return_ious): 115 | return collision_mask 116 | 117 | ret_value = [collision_mask,] 118 | if return_empty_grasp: 119 | inner_mask = (mask1 & mask2 & (~mask4) & (~mask6)) 120 | inner_volume = (heights * self.finger_length * widths / (self.voxel_size**3)).reshape(-1) 121 | empty_mask = (inner_mask.sum(axis=-1)/inner_volume < empty_thresh) 122 | ret_value.append(empty_mask) 123 | if return_ious: 124 | left_iou = left_mask.sum(axis=1) / (left_right_volume+1e-6) 125 | right_iou = right_mask.sum(axis=1) / (left_right_volume+1e-6) 126 | bottom_iou = bottom_mask.sum(axis=1) / (bottom_volume+1e-6) 127 | shifting_iou = shifting_mask.sum(axis=1) / (shifting_volume+1e-6) 128 | ret_value.append([global_iou, left_iou, right_iou, bottom_iou, shifting_iou]) 129 | return ret_value 130 | -------------------------------------------------------------------------------- /FGC_generate/vis_grasp.py: -------------------------------------------------------------------------------- 1 | from graspnetAPI import GraspNet 2 | import os 3 | import open3d as o3d 4 | import numpy as np 5 | from graspnetAPI.utils.utils import generate_views, get_model_grasps 6 | from graspnetAPI.utils.rotation import batch_viewpoint_params_to_matrix 7 | from graspnetAPI.grasp import Grasp 8 | 9 | 10 | def to_open3d_geometry_list(grip): 11 | ''' 12 | **Output:** 13 | 14 | - list of open3d.geometry.Geometry of the grippers. 15 | ''' 16 | geometry = [] 17 | 18 | for i in range(grip.shape[0]): 19 | if grip[i][0]==1: 20 | g = Grasp(grip[i]) 21 | geometry.append(g.to_open3d_geometry()) 22 | # g = Grasp(grip) 23 | # geometry.append(g.to_open3d_geometry()) 24 | 25 | return geometry 26 | 27 | 28 | def get_camera_parameters(camera='kinect'): 29 | ''' 30 | author: Minghao Gou 31 | 32 | **Input:** 33 | 34 | - camera: string of type of camera: 'kinect' or 'realsense' 35 | 36 | **Output:** 37 | 38 | - open3d.camera.PinholeCameraParameters 39 | ''' 40 | param = o3d.camera.PinholeCameraParameters() 41 | param.extrinsic = np.eye(4, dtype=np.float64) 42 | # param.intrinsic = o3d.camera.PinholeCameraIntrinsic() 43 | if camera == 'kinect': 44 | param.intrinsic.set_intrinsics(1280, 720, 631.5, 631.2, 639.5, 359.5) 45 | elif camera == 'realsense': 46 | param.intrinsic.set_intrinsics(1280, 720, 927.17, 927.37, 639.5, 359.5) 47 | return param 48 | 49 | 50 | def visObjGrasp(dataset_root, obj_idx, num_grasp=10, th=0.5, max_width=0.08, save_folder='save_fig', show=False): 51 | ''' 52 | Author: chenxi-wang 53 | 54 | **Input:** 55 | 56 | - dataset_root: str, graspnet dataset root 57 | 58 | - obj_idx: int, index of object model 59 | 60 | - num_grasp: int, number of sampled grasps 61 | 62 | - th: float, threshold of friction coefficient 63 | 64 | - max_width: float, only visualize grasps with width<=max_width 65 | 66 | - save_folder: str, folder to save screen captures 67 | 68 | - show: bool, show visualization in open3d window if set to True 69 | ''' 70 | plyfile = np.load(os.path.join(dataset_root, 'grasp_label', '000_labels.npz')) 71 | model = plyfile['points'] 72 | 73 | num_views, num_angles, num_depths = 300, 12, 4 74 | views = generate_views(num_views) 75 | 76 | # vis = o3d.visualization.Visualizer() 77 | # vis.create_window(width=1280, height=720) 78 | # ctr = vis.get_view_control() 79 | param = get_camera_parameters(camera='realsense') 80 | 81 | cam_pos = np.load(os.path.join(dataset_root, 'scenes', 'scene_0000', 'realsense', 'cam0_wrt_table.npy')) 82 | param.extrinsic = np.linalg.inv(cam_pos).tolist() 83 | 84 | sampled_points, offsets, scores, _ = get_model_grasps('%s/grasp_label/%03d_labels.npz' % (dataset_root, obj_idx)) 85 | # ----------------------------------------------------------------------------------------------------------------- 86 | point_inds = np.arange(sampled_points.shape[0]) 87 | template_views = generate_views(num_views) 88 | template_views = template_views[np.newaxis, :, np.newaxis, np.newaxis, :] 89 | template_views = np.tile(template_views, [1, 1, num_angles, num_depths, 1]) 90 | 91 | num_points = len(point_inds) 92 | target_points = sampled_points[:, np.newaxis, np.newaxis, np.newaxis, :] 93 | target_points = np.tile(target_points, [1, num_views, num_angles, num_depths, 1]) 94 | views = np.tile(template_views, [num_points, 1, 1, 1, 1]) 95 | angles = offsets[:, :, :, :, 0] 96 | depths = offsets[:, :, :, :, 1] 97 | widths = offsets[:, :, :, :, 2] 98 | 99 | mask1 = ((scores <= 0.4) & (scores > 0)) 100 | target_points = target_points[mask1] 101 | views = views[mask1] 102 | angles = angles[mask1] 103 | depths = depths[mask1] 104 | widths = widths[mask1] 105 | fric_coefs = scores[mask1] 106 | 107 | Rs = batch_viewpoint_params_to_matrix(-views, angles) 108 | 109 | num_grasp = widths.shape[0] 110 | scores = (1.1 - fric_coefs).reshape(-1, 1) 111 | widths = widths.reshape(-1, 1) 112 | heights = 0.02 * np.ones((num_grasp, 1)) 113 | depths = depths.reshape(-1, 1) 114 | rotations = Rs.reshape((-1, 9)) 115 | object_ids = obj_idx * np.ones((num_grasp, 1), dtype=np.int32) 116 | 117 | obj_grasp_array = np.hstack([scores, widths, heights, depths, rotations, target_points, object_ids]).astype( 118 | np.float32) 119 | 120 | return model, obj_grasp_array 121 | #cnt = 0 122 | # point_inds = np.arange(sampled_points.shape[0]) 123 | # np.random.shuffle(point_inds) 124 | # grippers = [] 125 | # 126 | # 127 | # for point_ind in point_inds: 128 | # target_point = sampled_points[point_ind] 129 | # offset = offsets[point_ind] 130 | # score = scores[point_ind] 131 | # view_inds = np.arange(300) 132 | # np.random.shuffle(view_inds) 133 | # flag = False 134 | # for v in view_inds: 135 | # if flag: break 136 | # view = views[v] 137 | # angle_inds = np.arange(12) 138 | # np.random.shuffle(angle_inds) 139 | # for a in angle_inds: 140 | # if flag: break 141 | # depth_inds = np.arange(4) 142 | # np.random.shuffle(depth_inds) 143 | # for d in depth_inds: 144 | # if flag: break 145 | # angle, depth, width = offset[v, a, d] 146 | # if score[v, a, d] > th or score[v, a, d] < 0 or width > max_width: 147 | # continue 148 | # R = viewpoint_params_to_matrix(-view, angle) 149 | # t = target_point 150 | # gripper = plot_gripper_pro_max(t, R, width, depth, 1.1 - score[v, a, d]) 151 | # grippers.append(gripper) 152 | # flag = True 153 | # if flag: 154 | # cnt += 1 155 | # if cnt == num_grasp: 156 | # break 157 | # 158 | # vis.add_geometry(model) 159 | # for gripper in grippers: 160 | # vis.add_geometry(gripper) 161 | # ctr.convert_from_pinhole_camera_parameters(param) 162 | # vis.poll_events() 163 | # filename = os.path.join(save_folder, 'object_{}_grasp.png'.format(obj_idx)) 164 | # vis.capture_screen_image(filename, do_render=True) 165 | # if show: 166 | # o3d.visualization.draw_geometries([model, *grippers]) 167 | 168 | 169 | if __name__=='__main__': 170 | root = '../grasp_data' 171 | obj_idx = 0 172 | model, grasp = visObjGrasp(root, obj_idx, show=True) 173 | modelpc = o3d.geometry.PointCloud() 174 | modelpc.points = o3d.utility.Vector3dVector(model) 175 | # gg = GraspGroup(grasp).nms(translation_thresh=0.5) 176 | # gg.sort_by_score() 177 | # gripper = gg.to_open3d_geometry_list() 178 | gripper = to_open3d_geometry_list(grasp) 179 | #np.random.shuffle(gripper) 180 | gripper_sample = gripper[:20] 181 | o3d.visualization.draw_geometries([modelpc, *gripper_sample]) 182 | 183 | -------------------------------------------------------------------------------- /FGC_generate/contact_score.py: -------------------------------------------------------------------------------- 1 | from vis_grasp import visObjGrasp, to_open3d_geometry_list 2 | import open3d as o3d 3 | import numpy as np 4 | import os 5 | from graspnetAPI.utils.utils import generate_views 6 | from graspnetAPI.utils.rotation import batch_viewpoint_params_to_matrix 7 | from tqdm import tqdm 8 | import torch 9 | 10 | 11 | class Contact_decision(): 12 | def __init__(self, date_root): 13 | self.date_root = date_root 14 | 15 | def load_grasp_label(self): 16 | obj_name = list(range(88)) 17 | obj_path = os.path.join(self.date_root, 'grasp_label', '{}_labels.npz'.format(str(0).zfill(3))) 18 | label = np.load(obj_path) 19 | points = label['points'] 20 | offsets = label['offsets'] 21 | scores = label['scores'] 22 | collision = label['collision'] 23 | obj_idx = 0 24 | return points, offsets, scores, collision, obj_idx 25 | 26 | def get_grasp(self): 27 | sampled_points, offsets, scores, _, obj_idx = self.load_grasp_label() 28 | num_views, num_angles, num_depths = 300, 12, 4 29 | 30 | point_inds = np.arange(sampled_points.shape[0]) 31 | template_views = generate_views(num_views) 32 | template_views = template_views[np.newaxis, :, np.newaxis, np.newaxis, :] 33 | template_views = np.tile(template_views, [1, 1, num_angles, num_depths, 1]) 34 | 35 | num_points = len(point_inds) 36 | target_points = sampled_points[:, np.newaxis, np.newaxis, np.newaxis, :] 37 | target_points = np.tile(target_points, [1, num_views, num_angles, num_depths, 1]) 38 | views = np.tile(template_views, [num_points, 1, 1, 1, 1]) 39 | angles = offsets[:, :, :, :, 0] 40 | depths = offsets[:, :, :, :, 1] 41 | widths = offsets[:, :, :, :, 2] 42 | 43 | mask1 = (scores > 0) 44 | mask1_idx = np.where(scores > 0) 45 | target_points = target_points[mask1] 46 | views = views[mask1] 47 | angles = angles[mask1] 48 | depths = depths[mask1] 49 | widths = widths[mask1] 50 | fric_coefs = scores[mask1] 51 | 52 | Rs = batch_viewpoint_params_to_matrix(-views, angles) 53 | 54 | num_grasp = widths.shape[0] 55 | scores = (1.1 - fric_coefs).reshape(-1, 1) 56 | widths = widths.reshape(-1, 1) 57 | heights = 0.02 * np.ones((num_grasp, 1)) 58 | depths = depths.reshape(-1, 1) 59 | rotations = Rs.reshape((-1, 9)) 60 | object_ids = obj_idx * np.ones((num_grasp, 1), dtype=np.int32) 61 | 62 | obj_grasp_array = np.hstack([scores, widths, heights, depths, rotations, target_points, object_ids]).astype( 63 | np.float32) 64 | 65 | return sampled_points, obj_grasp_array, mask1_idx 66 | 67 | def cal_dist(self, point1, point2, point3): 68 | ''' 69 | 70 | :param point1: (x1, y1, z1), the point 1 of line 71 | :param point2: (x2, y2, z2), the point 2 of line 72 | :param point3: (x3, y3, z3) 73 | v12 = point1-point2 74 | v13 = point1-point3 75 | distance = |v12×v13| / |v12| 76 | :return: dis 77 | 78 | ''' 79 | 80 | vec1 = point1-point2 81 | vec2 = point1-point3 82 | dis = abs(torch.linalg.norm(torch.cross(vec1, vec2)))/abs(torch.linalg.norm(vec1)) 83 | dis13_left = torch.linalg.norm(point1-point3) 84 | dis23_right = torch.linalg.norm(point2-point3) 85 | if dis13_left <= dis23_right: 86 | # 0 means point3 close to left contact, 1 means point3 close to right contact 87 | dis = [dis, 0] 88 | else: 89 | dis = [dis, 1] 90 | return dis 91 | 92 | def init_contact(self, width, depth): 93 | height = 0.004 94 | left_point = torch.tensor([depth - height / 2, -width / 2, 0]).cuda() # 定义抓取接触点的初始点 95 | right_point = torch.tensor([depth - height / 2, width / 2, 0]).cuda() 96 | return left_point, right_point 97 | 98 | def contact_find(self): 99 | obj_pc, grasp, mask1_idx = self.get_grasp() 100 | 101 | obj_pc_t = torch.from_numpy(obj_pc).cuda() 102 | grasp_t = torch.from_numpy(grasp).cuda() 103 | 104 | grasp_num = grasp_t.shape[0] 105 | #grasp = grasp[1] 106 | 107 | four_point_all = [] 108 | gravity_score_all = [] 109 | 110 | for x in tqdm(range(grasp_num), desc = 'Loading grasp...'): 111 | grasp_x = grasp_t[x] 112 | width = grasp_x[1] 113 | depth = grasp_x[3] 114 | rot = torch.reshape(grasp_x[4:13], (3, 3)) 115 | center = grasp_x[-4:-1].unsqueeze(1) # 3*1 116 | left_point, right_point = self.init_contact(width, depth) 117 | 118 | left_contact = torch.mm(rot, left_point.unsqueeze(1)) + center # 得到旋转平移后的接触点 3*1 119 | right_contact = torch.mm(rot, right_point.unsqueeze(1)) + center 120 | gravity_center = torch.tensor([0, 0, 0]).cuda() 121 | 122 | pc_num = obj_pc_t.shape[0] 123 | dis = torch.zeros((pc_num, 2)).cuda() 124 | for i in range(pc_num): 125 | point3 = obj_pc_t[i] 126 | dis_i = self.cal_dist(left_contact.squeeze(1), right_contact.squeeze(1), point3) 127 | dis[i, :] = torch.tensor(dis_i) 128 | 129 | min2max = torch.argsort(dis[:, 0]) 130 | for i in min2max: 131 | if dis[i, 1]==0: 132 | left_idx = i 133 | break 134 | 135 | for j in min2max: 136 | if dis[j, 1]==1: 137 | right_idx = j 138 | break 139 | 140 | point_target_left = obj_pc_t[left_idx] 141 | point_target_right = obj_pc_t[right_idx] 142 | gravity_center_score = self.cal_dist(point_target_left, point_target_right, gravity_center) 143 | 144 | four_point = torch.stack((point_target_left, point_target_right, left_contact.squeeze(1), right_contact.squeeze(1)), axis=0) 145 | 146 | four_point_all.append(four_point) 147 | gravity_score_all.append(gravity_center_score) 148 | four_point_all = torch.tensor(four_point_all) 149 | gravity_score_all = torch.tensor(gravity_score_all) 150 | return obj_pc_t, grasp, four_point_all, gravity_score_all 151 | 152 | def vis_contact(self): 153 | obj_pc, grasp, four_point, score_gc = self.contact_find() 154 | 155 | objp3d = o3d.geometry.PointCloud() 156 | objp3d.points = o3d.utility.Vector3dVector(obj_pc) 157 | objp3d.paint_uniform_color([0.3, 0.5, 0]) 158 | 159 | pc_target = o3d.geometry.PointCloud() 160 | pc_target.points = o3d.utility.Vector3dVector(four_point) 161 | pc_target.paint_uniform_color([0, 0, 1]) 162 | 163 | gg = to_open3d_geometry_list(grasp) 164 | 165 | o3d.visualization.draw_geometries([*gg, pc_target, objp3d], width=800, height=600, left=50, top=50) 166 | 167 | 168 | if __name__=='__main__': 169 | os.environ["CUDA_VISIBLE_DEVICES"] = '1' 170 | 171 | root = '../grasp_data' 172 | contact_decision = Contact_decision(root) 173 | x, y, z, w = contact_decision.contact_find() -------------------------------------------------------------------------------- /models/FGC_graspnet.py: -------------------------------------------------------------------------------- 1 | """ GraspNet baseline model definition. 2 | Author: chenxi-wang 3 | """ 4 | 5 | import os 6 | import sys 7 | import numpy as np 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | 12 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 13 | ROOT_DIR = os.path.dirname(BASE_DIR) 14 | sys.path.append(ROOT_DIR) 15 | sys.path.append(os.path.join(ROOT_DIR, 'pointnet2')) 16 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 17 | from utils.loss_utils import GRASP_MAX_WIDTH, GRASP_MAX_TOLERANCE, THRESH_GOOD, THRESH_BAD,\ 18 | transform_point_cloud, generate_grasp_views,\ 19 | batch_viewpoint_params_to_matrix, huber_loss 20 | import pytorch_utils as pt_utils 21 | from pointnet2.pointnet2_utils import CylinderQueryAndGroup 22 | from models.backbone import Pointnet2Backbone, Local_attention 23 | from models.modules import ApproachNet 24 | from utils.loss_utils import GRASP_MAX_WIDTH, GRASP_MAX_TOLERANCE 25 | from utils.label_generation import process_grasp_labels, match_grasp_view_and_label, batch_viewpoint_params_to_matrix 26 | 27 | 28 | 29 | class OperationNet(nn.Module): 30 | """ Grasp configure estimation. 31 | 32 | Input: 33 | num_angle: [int] 34 | number of in-plane rotation angle classes 35 | the value of the i-th class --> i*PI/num_angle (i=0,...,num_angle-1) 36 | num_depth: [int] 37 | number of gripper depth classes 38 | """ 39 | def __init__(self, num_angle, num_depth): 40 | # Output: 41 | super().__init__() 42 | self.num_angle = num_angle 43 | self.num_depth = num_depth 44 | 45 | self.conv1 = nn.Conv1d(256, 128, 1) 46 | self.conv2 = nn.Conv1d(128, 128, 1) 47 | self.conv3 = nn.Conv1d(128, num_angle*num_depth+num_angle+2*num_depth, 1) 48 | self.bn1 = nn.BatchNorm1d(128) 49 | self.bn2 = nn.BatchNorm1d(128) 50 | 51 | def forward(self, vp_features, end_points): 52 | """ Forward pass. 53 | 54 | Input: 55 | vp_features: [torch.FloatTensor, (batch_size,num_seed,3)] 56 | features of grouped points in different depths 57 | end_points: [dict] 58 | 59 | Output: 60 | end_points: [dict] 61 | """ 62 | B, D, N = vp_features.size() 63 | vp_features = F.relu(self.bn1(self.conv1(vp_features)), inplace=True) 64 | vp_features = F.relu(self.bn2(self.conv2(vp_features)), inplace=True) 65 | vp_features = self.conv3(vp_features) # B, D, N 66 | 67 | # split prediction 68 | end_points['grasp_score_pred'] = vp_features[:, 0:4*12] 69 | end_points['grasp_angle_cls_pred'] = vp_features[:, 48:48+self.num_angle] 70 | end_points['grasp_width_pred'] = vp_features[:, 48+self.num_angle:48+self.num_angle+self.num_depth] 71 | end_points['grasp_depth_cls_pred'] = vp_features[:, 48+self.num_angle+self.num_depth:] 72 | return end_points 73 | 74 | 75 | class CloudCrop(nn.Module): 76 | """ Cylinder group and align for grasp configure estimation. Return a list of grouped points with different cropping depths. 77 | 78 | Input: 79 | nsample: [int] 80 | sample number in a group 81 | seed_feature_dim: [int] 82 | number of channels of grouped points 83 | cylinder_radius: [float] 84 | radius of the cylinder space 85 | hmin: [float] 86 | height of the bottom surface 87 | hmax_list: [list of float] 88 | list of heights of the upper surface 89 | """ 90 | 91 | def __init__(self, nsample, seed_feature_dim, cylinder_radius=0.05, hmin=-0.02, hmax=0.02): 92 | super().__init__() 93 | self.nsample = nsample # 64 94 | self.in_dim = seed_feature_dim # 3 95 | self.cylinder_radius = cylinder_radius 96 | mlps = [128+3, 256] 97 | 98 | self.groupers = CylinderQueryAndGroup( 99 | cylinder_radius, hmin, hmax, nsample, use_xyz=True) 100 | self.mlps = pt_utils.SharedMLP(mlps, bn=True) 101 | self.local_att = Local_attention(256) 102 | 103 | def forward(self, seed_xyz, pointcloud, vp_rot, up_feature): 104 | """ Forward pass. 105 | 106 | Input: 107 | seed_xyz: [torch.FloatTensor, (batch_size,num_seed,3)] 108 | coordinates of seed points 109 | pointcloud: [torch.FloatTensor, (batch_size,num_seed,3)] 110 | the points to be cropped 111 | vp_rot: [torch.FloatTensor, (batch_size,num_seed,3,3)] 112 | rotation matrices generated from approach vectors 113 | 114 | Output: 115 | vp_features: [torch.FloatTensor, (batch_size,num_features,num_seed,num_depth)] 116 | features of grouped points in different depths 117 | """ 118 | B, num_seed, _, _ = vp_rot.size() 119 | grouped_features = self.groupers( 120 | pointcloud, seed_xyz, vp_rot, features=up_feature) # (batch_size, feature_dim, nsample) 121 | vp_features = self.mlps( 122 | grouped_features) # (batch_size, mlps[-1], num_seed, nsample) 123 | 124 | vp_features = vp_features.permute(0, 2, 1, 3).contiguous().view(B * num_seed, 256, 125 | self.nsample) # (B*num_seed*num_depth, C, K) 126 | vp_features = self.local_att(vp_features).contiguous().view(B, num_seed, 256, self.nsample).permute( 127 | 0, 2, 1, 3) 128 | 129 | vp_features = F.max_pool2d( 130 | vp_features, kernel_size=[1, vp_features.size(3)]).squeeze(-1) # (batch_size, mlps[-1], num_seed) 131 | return vp_features 132 | 133 | 134 | class FGC_graspnet(nn.Module): 135 | def __init__(self, input_feature_dim=0, num_view=300, num_angle=12, num_depth=4, cylinder_radius=0.05, hmin=-0.02, 136 | hmax=0.02, is_training=True, is_demo=False): 137 | super().__init__() 138 | self.backbone = Pointnet2Backbone(input_feature_dim) 139 | 140 | self.vpmodule = ApproachNet(num_view, 256) 141 | 142 | self.operation = OperationNet(num_angle, num_depth) 143 | self.local_att = CloudCrop(32, 3, cylinder_radius, hmin, hmax) 144 | self.num_depth = num_depth 145 | self.is_traning = is_training 146 | self.is_demo = is_demo 147 | 148 | def forward(self, end_points): 149 | pointcloud = end_points['point_clouds'] 150 | seed_features, seed_xyz, end_points = self.backbone(pointcloud, end_points) 151 | end_points = self.vpmodule(seed_xyz, seed_features, end_points) 152 | 153 | if not self.is_demo: 154 | end_points = process_grasp_labels(end_points) 155 | 156 | if self.is_traning: 157 | grasp_top_views_rot, _, _, end_points = match_grasp_view_and_label(end_points) 158 | seed_xyz = end_points['batch_grasp_point'] 159 | else: 160 | # _, _, _, end_points = match_grasp_view_and_label(end_points) 161 | grasp_top_views_rot = end_points['grasp_top_view_rot'] 162 | seed_xyz = end_points['fp2_xyz'] 163 | 164 | up_features = end_points['sa1_features'] # B, 128, 1024*4 165 | xyz = end_points['sa1_xyz'] 166 | vp_features = self.local_att(seed_xyz, xyz, grasp_top_views_rot, up_features) 167 | 168 | #vp_features = seed_features.permute(0, 2, 1).repeat(1, self.num_depth, 1) 169 | end_points = self.operation(vp_features, end_points) 170 | return end_points 171 | 172 | -------------------------------------------------------------------------------- /utils/label_generation.py: -------------------------------------------------------------------------------- 1 | """ Dynamically generate grasp labels during training. 2 | Author: chenxi-wang 3 | """ 4 | 5 | import os 6 | import sys 7 | import torch 8 | 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | ROOT_DIR = os.path.dirname(BASE_DIR) 11 | sys.path.append(ROOT_DIR) 12 | sys.path.append(os.path.join(ROOT_DIR, 'knn')) 13 | 14 | from knn_modules import knn 15 | from loss_utils import GRASP_MAX_WIDTH, batch_viewpoint_params_to_matrix,\ 16 | transform_point_cloud, generate_grasp_views 17 | 18 | def process_grasp_labels(end_points): 19 | """ Process labels according to scene points and object poses. """ 20 | clouds = end_points['input_xyz'] #(B, N, 3) 21 | seed_xyzs = end_points['fp2_xyz'] #(B, Ns, 3) change_here 22 | batch_size, num_samples, _ = seed_xyzs.size() 23 | 24 | batch_grasp_points = [] 25 | batch_grasp_views = [] 26 | batch_grasp_views_rot = [] 27 | batch_grasp_labels = [] 28 | batch_grasp_offsets = [] 29 | batch_grasp_tolerance = [] 30 | for i in range(len(clouds)): 31 | seed_xyz = seed_xyzs[i] #(Ns, 3) 32 | poses = end_points['object_poses_list'][i] #[(3, 4),] 33 | 34 | # get merged grasp points for label computation 35 | grasp_points_merged = [] 36 | grasp_views_merged = [] 37 | grasp_views_rot_merged = [] 38 | grasp_labels_merged = [] 39 | grasp_offsets_merged = [] 40 | for obj_idx, pose in enumerate(poses): 41 | grasp_points = end_points['grasp_points_list'][i][obj_idx] #(Np, 3) 42 | grasp_labels = end_points['grasp_labels_list'][i][obj_idx] #(Np, V, A, D) 43 | grasp_offsets = end_points['grasp_offsets_list'][i][obj_idx] #(Np, V, A, D, 3) 44 | 45 | _, V, A, D = grasp_labels.size() 46 | num_grasp_points = grasp_points.size(0) 47 | # generate and transform template grasp views 48 | grasp_views = generate_grasp_views(V).to(pose.device) #(V, 3) 49 | grasp_points_trans = transform_point_cloud(grasp_points, pose, '3x4') 50 | grasp_views_trans = transform_point_cloud(grasp_views, pose[:3,:3], '3x3') 51 | # generate and transform template grasp view rotation 52 | angles = torch.zeros(grasp_views.size(0), dtype=grasp_views.dtype, device=grasp_views.device) 53 | grasp_views_rot = batch_viewpoint_params_to_matrix(-grasp_views, angles) #(V, 3, 3) 54 | grasp_views_rot_trans = torch.matmul(pose[:3,:3], grasp_views_rot) #(V, 3, 3) 55 | 56 | # assign views 57 | grasp_views_ = grasp_views.transpose(0, 1).contiguous().unsqueeze(0) 58 | grasp_views_trans_ = grasp_views_trans.transpose(0, 1).contiguous().unsqueeze(0) 59 | view_inds = knn(grasp_views_trans_, grasp_views_, k=1).squeeze() - 1 60 | grasp_views_trans = torch.index_select(grasp_views_trans, 0, view_inds) #(V, 3) 61 | grasp_views_trans = grasp_views_trans.unsqueeze(0).expand(num_grasp_points, -1, -1) #(Np, V, 3) 62 | grasp_views_rot_trans = torch.index_select(grasp_views_rot_trans, 0, view_inds) #(V, 3, 3) 63 | grasp_views_rot_trans = grasp_views_rot_trans.unsqueeze(0).expand(num_grasp_points, -1, -1, -1) #(Np, V, 3, 3) 64 | grasp_labels = torch.index_select(grasp_labels, 1, view_inds) #(Np, V, A, D) 65 | grasp_offsets = torch.index_select(grasp_offsets, 1, view_inds) #(Np, V, A, D, 3) 66 | 67 | # add to list 68 | grasp_points_merged.append(grasp_points_trans) 69 | grasp_views_merged.append(grasp_views_trans) 70 | grasp_views_rot_merged.append(grasp_views_rot_trans) 71 | grasp_labels_merged.append(grasp_labels) 72 | grasp_offsets_merged.append(grasp_offsets) 73 | 74 | grasp_points_merged = torch.cat(grasp_points_merged, dim=0) #(Np', 3) 75 | grasp_views_merged = torch.cat(grasp_views_merged, dim=0) #(Np', V, 3) 76 | grasp_views_rot_merged = torch.cat(grasp_views_rot_merged, dim=0) #(Np', V, 3, 3) 77 | grasp_labels_merged = torch.cat(grasp_labels_merged, dim=0) #(Np', V, A, D) 78 | grasp_offsets_merged = torch.cat(grasp_offsets_merged, dim=0) #(Np', V, A, D, 3) 79 | 80 | # compute nearest neighbors 81 | seed_xyz_ = seed_xyz.transpose(0, 1).contiguous().unsqueeze(0) #(1, 3, Ns) 82 | grasp_points_merged_ = grasp_points_merged.transpose(0, 1).contiguous().unsqueeze(0) #(1, 3, Np') 83 | nn_inds = knn(grasp_points_merged_, seed_xyz_, k=1).squeeze() - 1 #(Ns) 84 | 85 | # assign anchor points to real points 86 | grasp_points_merged = torch.index_select(grasp_points_merged, 0, nn_inds) # (Ns, 3) 87 | grasp_views_merged = torch.index_select(grasp_views_merged, 0, nn_inds) # (Ns, V, 3) 88 | grasp_views_rot_merged = torch.index_select(grasp_views_rot_merged, 0, nn_inds) #(Ns, V, 3, 3) 89 | grasp_labels_merged = torch.index_select(grasp_labels_merged, 0, nn_inds) # (Ns, V, A, D) 90 | grasp_offsets_merged = torch.index_select(grasp_offsets_merged, 0, nn_inds) # (Ns, V, A, D, 3) 91 | 92 | # add to batch 93 | batch_grasp_points.append(grasp_points_merged) 94 | batch_grasp_views.append(grasp_views_merged) 95 | batch_grasp_views_rot.append(grasp_views_rot_merged) 96 | batch_grasp_labels.append(grasp_labels_merged) 97 | batch_grasp_offsets.append(grasp_offsets_merged) 98 | 99 | batch_grasp_points = torch.stack(batch_grasp_points, 0) #(B, Ns, 3) 100 | batch_grasp_views = torch.stack(batch_grasp_views, 0) #(B, Ns, V, 3) 101 | batch_grasp_views_rot = torch.stack(batch_grasp_views_rot, 0) #(B, Ns, V, 3, 3) 102 | batch_grasp_labels = torch.stack(batch_grasp_labels, 0) #(B, Ns, V, A, D) 103 | batch_grasp_offsets = torch.stack(batch_grasp_offsets, 0) #(B, Ns, V, A, D, 3) 104 | 105 | # process labels 106 | batch_grasp_widths = batch_grasp_offsets[:,:,:,:,:,2] 107 | label_mask = (batch_grasp_labels > 0) & (batch_grasp_widths <= GRASP_MAX_WIDTH) 108 | u_max = batch_grasp_labels.max() 109 | #batch_grasp_labels[label_mask] = torch.log(u_max / batch_grasp_labels[label_mask]) 110 | batch_grasp_labels[~label_mask] = 0 111 | a_max = batch_grasp_labels.max() 112 | batch_grasp_view_scores, _ = batch_grasp_labels.view(batch_size, num_samples, V, A*D).max(dim=-1) 113 | 114 | end_points['batch_grasp_point'] = batch_grasp_points 115 | end_points['batch_grasp_view'] = batch_grasp_views 116 | end_points['batch_grasp_view_rot'] = batch_grasp_views_rot 117 | end_points['batch_grasp_label'] = batch_grasp_labels 118 | end_points['batch_grasp_offset'] = batch_grasp_offsets 119 | end_points['batch_grasp_view_label'] = batch_grasp_view_scores.float() 120 | 121 | return end_points 122 | 123 | def match_grasp_view_and_label(end_points): 124 | """ Slice grasp labels according to predicted views. """ 125 | top_view_inds = end_points['grasp_top_view_inds'] # (B, Ns) 126 | template_views_rot = end_points['batch_grasp_view_rot'] # (B, Ns, V, 3, 3) 127 | grasp_labels = end_points['batch_grasp_label'] # (B, Ns, V, A, D) 128 | grasp_offsets = end_points['batch_grasp_offset'] # (B, Ns, V, A, D, 3) 129 | 130 | B, Ns, V, A, D = grasp_labels.size() 131 | top_view_inds_ = top_view_inds.view(B, Ns, 1, 1, 1).expand(-1, -1, -1, 3, 3) 132 | top_template_views_rot = torch.gather(template_views_rot, 2, top_view_inds_).squeeze(2) 133 | top_view_inds_ = top_view_inds.view(B, Ns, 1, 1, 1).expand(-1, -1, -1, A, D) 134 | top_view_grasp_labels = torch.gather(grasp_labels, 2, top_view_inds_).squeeze(2) 135 | top_view_inds_ = top_view_inds.view(B, Ns, 1, 1, 1, 1).expand(-1, -1, -1, A, D, 3) 136 | top_view_grasp_offsets = torch.gather(grasp_offsets, 2, top_view_inds_).squeeze(2) 137 | 138 | end_points['batch_grasp_view_rot'] = top_template_views_rot 139 | end_points['batch_grasp_label'] = top_view_grasp_labels # (B, Ns, A, D) 140 | end_points['batch_grasp_offset'] = top_view_grasp_offsets 141 | 142 | return top_template_views_rot, top_view_grasp_labels, top_view_grasp_offsets, end_points -------------------------------------------------------------------------------- /pointnet2/_ext_src/src/sampling_gpu.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | // 3 | // This source code is licensed under the MIT license found in the 4 | // LICENSE file in the root directory of this source tree. 5 | 6 | #include 7 | #include 8 | 9 | #include "cuda_utils.h" 10 | 11 | // input: points(b, c, n) idx(b, m) 12 | // output: out(b, c, m) 13 | __global__ void gather_points_kernel(int b, int c, int n, int m, 14 | const float *__restrict__ points, 15 | const int *__restrict__ idx, 16 | float *__restrict__ out) { 17 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 18 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 19 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 20 | int a = idx[i * m + j]; 21 | out[(i * c + l) * m + j] = points[(i * c + l) * n + a]; 22 | } 23 | } 24 | } 25 | } 26 | 27 | void gather_points_kernel_wrapper(int b, int c, int n, int npoints, 28 | const float *points, const int *idx, 29 | float *out) { 30 | gather_points_kernel<<>>(b, c, n, npoints, 32 | points, idx, out); 33 | 34 | CUDA_CHECK_ERRORS(); 35 | } 36 | 37 | // input: grad_out(b, c, m) idx(b, m) 38 | // output: grad_points(b, c, n) 39 | __global__ void gather_points_grad_kernel(int b, int c, int n, int m, 40 | const float *__restrict__ grad_out, 41 | const int *__restrict__ idx, 42 | float *__restrict__ grad_points) { 43 | for (int i = blockIdx.x; i < b; i += gridDim.x) { 44 | for (int l = blockIdx.y; l < c; l += gridDim.y) { 45 | for (int j = threadIdx.x; j < m; j += blockDim.x) { 46 | int a = idx[i * m + j]; 47 | atomicAdd(grad_points + (i * c + l) * n + a, 48 | grad_out[(i * c + l) * m + j]); 49 | } 50 | } 51 | } 52 | } 53 | 54 | void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints, 55 | const float *grad_out, const int *idx, 56 | float *grad_points) { 57 | gather_points_grad_kernel<<>>( 59 | b, c, n, npoints, grad_out, idx, grad_points); 60 | 61 | CUDA_CHECK_ERRORS(); 62 | } 63 | 64 | __device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i, 65 | int idx1, int idx2) { 66 | const float v1 = dists[idx1], v2 = dists[idx2]; 67 | const int i1 = dists_i[idx1], i2 = dists_i[idx2]; 68 | dists[idx1] = max(v1, v2); 69 | dists_i[idx1] = v2 > v1 ? i2 : i1; 70 | } 71 | 72 | // Input dataset: (b, n, 3), tmp: (b, n) 73 | // Ouput idxs (b, m) 74 | template 75 | __global__ void furthest_point_sampling_kernel( 76 | int b, int n, int m, const float *__restrict__ dataset, 77 | float *__restrict__ temp, int *__restrict__ idxs) { 78 | if (m <= 0) return; 79 | __shared__ float dists[block_size]; 80 | __shared__ int dists_i[block_size]; 81 | 82 | int batch_index = blockIdx.x; 83 | dataset += batch_index * n * 3; 84 | temp += batch_index * n; 85 | idxs += batch_index * m; 86 | 87 | int tid = threadIdx.x; 88 | const int stride = block_size; 89 | 90 | int old = 0; 91 | if (threadIdx.x == 0) idxs[0] = old; 92 | 93 | __syncthreads(); 94 | for (int j = 1; j < m; j++) { 95 | int besti = 0; 96 | float best = -1; 97 | float x1 = dataset[old * 3 + 0]; 98 | float y1 = dataset[old * 3 + 1]; 99 | float z1 = dataset[old * 3 + 2]; 100 | for (int k = tid; k < n; k += stride) { 101 | float x2, y2, z2; 102 | x2 = dataset[k * 3 + 0]; 103 | y2 = dataset[k * 3 + 1]; 104 | z2 = dataset[k * 3 + 2]; 105 | float mag = (x2 * x2) + (y2 * y2) + (z2 * z2); 106 | if (mag <= 1e-3) continue; 107 | 108 | float d = 109 | (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); 110 | 111 | float d2 = min(d, temp[k]); 112 | temp[k] = d2; 113 | besti = d2 > best ? k : besti; 114 | best = d2 > best ? d2 : best; 115 | } 116 | dists[tid] = best; 117 | dists_i[tid] = besti; 118 | __syncthreads(); 119 | 120 | if (block_size >= 512) { 121 | if (tid < 256) { 122 | __update(dists, dists_i, tid, tid + 256); 123 | } 124 | __syncthreads(); 125 | } 126 | if (block_size >= 256) { 127 | if (tid < 128) { 128 | __update(dists, dists_i, tid, tid + 128); 129 | } 130 | __syncthreads(); 131 | } 132 | if (block_size >= 128) { 133 | if (tid < 64) { 134 | __update(dists, dists_i, tid, tid + 64); 135 | } 136 | __syncthreads(); 137 | } 138 | if (block_size >= 64) { 139 | if (tid < 32) { 140 | __update(dists, dists_i, tid, tid + 32); 141 | } 142 | __syncthreads(); 143 | } 144 | if (block_size >= 32) { 145 | if (tid < 16) { 146 | __update(dists, dists_i, tid, tid + 16); 147 | } 148 | __syncthreads(); 149 | } 150 | if (block_size >= 16) { 151 | if (tid < 8) { 152 | __update(dists, dists_i, tid, tid + 8); 153 | } 154 | __syncthreads(); 155 | } 156 | if (block_size >= 8) { 157 | if (tid < 4) { 158 | __update(dists, dists_i, tid, tid + 4); 159 | } 160 | __syncthreads(); 161 | } 162 | if (block_size >= 4) { 163 | if (tid < 2) { 164 | __update(dists, dists_i, tid, tid + 2); 165 | } 166 | __syncthreads(); 167 | } 168 | if (block_size >= 2) { 169 | if (tid < 1) { 170 | __update(dists, dists_i, tid, tid + 1); 171 | } 172 | __syncthreads(); 173 | } 174 | 175 | old = dists_i[0]; 176 | if (tid == 0) idxs[j] = old; 177 | } 178 | } 179 | 180 | void furthest_point_sampling_kernel_wrapper(int b, int n, int m, 181 | const float *dataset, float *temp, 182 | int *idxs) { 183 | unsigned int n_threads = opt_n_threads(n); 184 | 185 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 186 | 187 | switch (n_threads) { 188 | case 512: 189 | furthest_point_sampling_kernel<512> 190 | <<>>(b, n, m, dataset, temp, idxs); 191 | break; 192 | case 256: 193 | furthest_point_sampling_kernel<256> 194 | <<>>(b, n, m, dataset, temp, idxs); 195 | break; 196 | case 128: 197 | furthest_point_sampling_kernel<128> 198 | <<>>(b, n, m, dataset, temp, idxs); 199 | break; 200 | case 64: 201 | furthest_point_sampling_kernel<64> 202 | <<>>(b, n, m, dataset, temp, idxs); 203 | break; 204 | case 32: 205 | furthest_point_sampling_kernel<32> 206 | <<>>(b, n, m, dataset, temp, idxs); 207 | break; 208 | case 16: 209 | furthest_point_sampling_kernel<16> 210 | <<>>(b, n, m, dataset, temp, idxs); 211 | break; 212 | case 8: 213 | furthest_point_sampling_kernel<8> 214 | <<>>(b, n, m, dataset, temp, idxs); 215 | break; 216 | case 4: 217 | furthest_point_sampling_kernel<4> 218 | <<>>(b, n, m, dataset, temp, idxs); 219 | break; 220 | case 2: 221 | furthest_point_sampling_kernel<2> 222 | <<>>(b, n, m, dataset, temp, idxs); 223 | break; 224 | case 1: 225 | furthest_point_sampling_kernel<1> 226 | <<>>(b, n, m, dataset, temp, idxs); 227 | break; 228 | default: 229 | furthest_point_sampling_kernel<512> 230 | <<>>(b, n, m, dataset, temp, idxs); 231 | } 232 | 233 | CUDA_CHECK_ERRORS(); 234 | } 235 | -------------------------------------------------------------------------------- /FGC_generate/contact_score_numba.py: -------------------------------------------------------------------------------- 1 | from vis_grasp import visObjGrasp, to_open3d_geometry_list 2 | import open3d as o3d 3 | import numpy as np 4 | import os 5 | from graspnetAPI.utils.utils import generate_views 6 | from graspnetAPI.utils.rotation import batch_viewpoint_params_to_matrix 7 | from tqdm import tqdm 8 | from numba import jit 9 | 10 | 11 | @jit(nopython=True) 12 | class Contact_decision(): 13 | def __init__(self, date_root): 14 | self.date_root = date_root 15 | 16 | def load_grasp_label(self): 17 | obj_name = list(range(88)) 18 | obj_path = os.path.join(self.date_root, 'grasp_label', '{}_labels.npz'.format(str(0).zfill(3))) 19 | label = np.load(obj_path) 20 | points = label['points'] 21 | offsets = label['offsets'] 22 | scores = label['scores'] 23 | collision = label['collision'] 24 | obj_idx = 0 25 | return points, offsets, scores, collision, obj_idx 26 | 27 | def get_grasp(self): 28 | sampled_points, offsets, scores, _, obj_idx = self.load_grasp_label() 29 | num_views, num_angles, num_depths = 300, 12, 4 30 | 31 | point_inds = np.arange(sampled_points.shape[0]) 32 | template_views = generate_views(num_views) 33 | template_views = template_views[np.newaxis, :, np.newaxis, np.newaxis, :] 34 | template_views = np.tile(template_views, [1, 1, num_angles, num_depths, 1]) 35 | 36 | num_points = len(point_inds) 37 | target_points = sampled_points[:, np.newaxis, np.newaxis, np.newaxis, :] 38 | target_points = np.tile(target_points, [1, num_views, num_angles, num_depths, 1]) 39 | views = np.tile(template_views, [num_points, 1, 1, 1, 1]) 40 | angles = offsets[:, :, :, :, 0] 41 | depths = offsets[:, :, :, :, 1] 42 | widths = offsets[:, :, :, :, 2] 43 | 44 | mask1 = (scores > 0) 45 | mask1_idx = np.where(scores > 0) 46 | target_points = target_points[mask1] 47 | views = views[mask1] 48 | angles = angles[mask1] 49 | depths = depths[mask1] 50 | widths = widths[mask1] 51 | fric_coefs = scores[mask1] 52 | 53 | Rs = batch_viewpoint_params_to_matrix(-views, angles) 54 | 55 | num_grasp = widths.shape[0] 56 | scores = (1.1 - fric_coefs).reshape(-1, 1) 57 | widths = widths.reshape(-1, 1) 58 | heights = 0.02 * np.ones((num_grasp, 1)) 59 | depths = depths.reshape(-1, 1) 60 | rotations = Rs.reshape((-1, 9)) 61 | object_ids = obj_idx * np.ones((num_grasp, 1), dtype=np.int32) 62 | 63 | obj_grasp_array = np.hstack([scores, widths, heights, depths, rotations, target_points, object_ids]).astype( 64 | np.float32) 65 | 66 | return sampled_points, obj_grasp_array, mask1_idx 67 | 68 | def cal_dist(self, point1, point2, point3): 69 | ''' 70 | 71 | :param point1: (x1, y1, z1), the point 1 of line 72 | :param point2: (x2, y2, z2), the point 2 of line 73 | :param point3: (x3, y3, z3) 74 | v12 = point1-point2 75 | v13 = point1-point3 76 | distance = |v12×v13| / |v12| 77 | :return: dis 78 | 79 | ''' 80 | 81 | vec1 = point1-point2 82 | vec2 = point1-point3 83 | dis = abs(np.linalg.norm(np.cross(vec1, vec2)))/abs(np.linalg.norm(vec1)) 84 | dis13_left = np.linalg.norm(point1-point3) 85 | dis23_right = np.linalg.norm(point2-point3) 86 | if dis13_left <= dis23_right: 87 | # 0 means point3 close to left contact, 1 means point3 close to right contact 88 | dis = [dis, 0] 89 | else: 90 | dis = [dis, 1] 91 | return dis 92 | 93 | def init_contact(self, width, depth): 94 | height = 0.004 95 | left_point = np.array([depth - height / 2, -width / 2, 0]) # 定义抓取接触点的初始点 96 | right_point = np.array([depth - height / 2, width / 2, 0]) 97 | return left_point, right_point 98 | 99 | def contact_find_solo(self): 100 | obj_pc, grasp, mask1_idx = self.get_grasp() 101 | 102 | grasp = grasp[1] 103 | 104 | width = grasp[1] 105 | depth = grasp[3] 106 | rot = grasp[4:13].reshape((3, 3)) 107 | center = grasp[-4:-1] 108 | left_point, right_point = self.init_contact(width, depth) 109 | 110 | left_contact = np.dot(rot, left_point.T).T + center # 得到旋转平移后的接触点 111 | right_contact = np.dot(rot, right_point.T).T + center 112 | gravity_center = [0, 0, 0] 113 | 114 | pc_num = obj_pc.shape[0] 115 | dis = np.zeros((pc_num, 2)) 116 | for i in range(pc_num): 117 | point3 = obj_pc[i] 118 | dis_i = self.cal_dist(left_contact, right_contact, point3) 119 | dis[i, :] = np.asarray(dis_i) 120 | 121 | min2max = np.argsort(dis[:, 0]) 122 | for i in min2max: 123 | if dis[i, 1] == 0: 124 | left_idx = i 125 | break 126 | 127 | for j in min2max: 128 | if dis[j, 1] == 1: 129 | right_idx = j 130 | break 131 | 132 | point_target_left = obj_pc[left_idx] 133 | point_target_right = obj_pc[right_idx] 134 | gravity_center_score = self.cal_dist(point_target_left, point_target_right, gravity_center) 135 | 136 | four_point = np.stack((point_target_left, point_target_right, left_contact, right_contact), axis=0) 137 | 138 | return obj_pc, grasp, four_point, gravity_center_score 139 | 140 | def contact_find(self): 141 | obj_pc, grasp, mask1_idx = self.get_grasp() 142 | 143 | grasp_num = grasp.shape[0] 144 | #grasp = grasp[1] 145 | 146 | four_point_all = [] 147 | gravity_score_all = [] 148 | 149 | for x in tqdm(range(grasp_num), desc = 'Loading grasp...'): 150 | grasp_x = grasp[x] 151 | width = grasp_x[1] 152 | depth = grasp_x[3] 153 | rot = grasp_x[4:13].reshape((3, 3)) 154 | center = grasp_x[-4:-1] 155 | left_point, right_point = self.init_contact(width, depth) 156 | 157 | left_contact = np.dot(rot, left_point.T).T + center # 得到旋转平移后的接触点 158 | right_contact = np.dot(rot, right_point.T).T + center 159 | gravity_center = [0, 0, 0] 160 | 161 | pc_num = obj_pc.shape[0] 162 | dis = np.zeros((pc_num, 2)) 163 | for i in range(pc_num): 164 | point3 = obj_pc[i] 165 | dis_i = self.cal_dist(left_contact, right_contact, point3) 166 | dis[i, :] = np.asarray(dis_i) 167 | 168 | min2max = np.argsort(dis[:, 0]) 169 | for i in min2max: 170 | if dis[i, 1]==0: 171 | left_idx = i 172 | break 173 | 174 | for j in min2max: 175 | if dis[j, 1]==1: 176 | right_idx = j 177 | break 178 | 179 | point_target_left = obj_pc[left_idx] 180 | point_target_right = obj_pc[right_idx] 181 | gravity_center_score = self.cal_dist(point_target_left, point_target_right, gravity_center) 182 | 183 | four_point = np.stack((point_target_left, point_target_right, left_contact, right_contact), axis=0) 184 | 185 | from rnn_neighbor import Curvity_score 186 | curve_score = Curvity_score(obj_pc, point_target_left, point_target_right) 187 | 188 | four_point_all.append(four_point) 189 | gravity_score_all.append(gravity_score_all) 190 | four_point_all = np.asarray(four_point_all) 191 | gravity_score_all = np.asarray(gravity_score_all) 192 | return obj_pc, grasp, four_point_all, gravity_score_all 193 | 194 | def vis_contact(self): 195 | obj_pc, grasp, four_point, score_gc = self.contact_find_solo() 196 | 197 | objp3d = o3d.geometry.PointCloud() 198 | objp3d.points = o3d.utility.Vector3dVector(obj_pc) 199 | objp3d.paint_uniform_color([0.3, 0.5, 0]) 200 | 201 | pc_target = o3d.geometry.PointCloud() 202 | pc_target.points = o3d.utility.Vector3dVector(four_point) 203 | pc_target.paint_uniform_color([0, 0, 1]) 204 | 205 | gg = to_open3d_geometry_list(grasp) 206 | 207 | o3d.visualization.draw_geometries([*gg, pc_target, objp3d], width=800, height=600, left=50, top=50) 208 | 209 | 210 | if __name__=='__main__': 211 | root = '../grasp_data' 212 | contact_decision = Contact_decision(root) 213 | #x, y, z, w = contact_decision.contact_find() 214 | contact_decision.vis_contact() -------------------------------------------------------------------------------- /pointnet2/pytorch_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | # 3 | # This source code is licensed under the MIT license found in the 4 | # LICENSE file in the root directory of this source tree. 5 | 6 | ''' Modified based on Ref: https://github.com/erikwijmans/Pointnet2_PyTorch ''' 7 | import torch 8 | import torch.nn as nn 9 | from typing import List, Tuple 10 | 11 | class SharedMLP(nn.Sequential): 12 | 13 | def __init__( 14 | self, 15 | args: List[int], 16 | *, 17 | bn: bool = False, 18 | activation=nn.ReLU(inplace=True), 19 | preact: bool = False, 20 | first: bool = False, 21 | name: str = "" 22 | ): 23 | super().__init__() 24 | 25 | for i in range(len(args) - 1): 26 | self.add_module( 27 | name + 'layer{}'.format(i), 28 | Conv2d( 29 | args[i], 30 | args[i + 1], 31 | bn=(not first or not preact or (i != 0)) and bn, 32 | activation=activation 33 | if (not first or not preact or (i != 0)) else None, 34 | preact=preact 35 | ) 36 | ) 37 | 38 | 39 | class _BNBase(nn.Sequential): 40 | 41 | def __init__(self, in_size, batch_norm=None, name=""): 42 | super().__init__() 43 | self.add_module(name + "bn", batch_norm(in_size)) 44 | 45 | nn.init.constant_(self[0].weight, 1.0) 46 | nn.init.constant_(self[0].bias, 0) 47 | 48 | 49 | class BatchNorm1d(_BNBase): 50 | 51 | def __init__(self, in_size: int, *, name: str = ""): 52 | super().__init__(in_size, batch_norm=nn.BatchNorm1d, name=name) 53 | 54 | 55 | class BatchNorm2d(_BNBase): 56 | 57 | def __init__(self, in_size: int, name: str = ""): 58 | super().__init__(in_size, batch_norm=nn.BatchNorm2d, name=name) 59 | 60 | 61 | class BatchNorm3d(_BNBase): 62 | 63 | def __init__(self, in_size: int, name: str = ""): 64 | super().__init__(in_size, batch_norm=nn.BatchNorm3d, name=name) 65 | 66 | 67 | class _ConvBase(nn.Sequential): 68 | 69 | def __init__( 70 | self, 71 | in_size, 72 | out_size, 73 | kernel_size, 74 | stride, 75 | padding, 76 | activation, 77 | bn, 78 | init, 79 | conv=None, 80 | batch_norm=None, 81 | bias=True, 82 | preact=False, 83 | name="" 84 | ): 85 | super().__init__() 86 | 87 | bias = bias and (not bn) 88 | conv_unit = conv( 89 | in_size, 90 | out_size, 91 | kernel_size=kernel_size, 92 | stride=stride, 93 | padding=padding, 94 | bias=bias 95 | ) 96 | init(conv_unit.weight) 97 | if bias: 98 | nn.init.constant_(conv_unit.bias, 0) 99 | 100 | if bn: 101 | if not preact: 102 | bn_unit = batch_norm(out_size) 103 | else: 104 | bn_unit = batch_norm(in_size) 105 | 106 | if preact: 107 | if bn: 108 | self.add_module(name + 'bn', bn_unit) 109 | 110 | if activation is not None: 111 | self.add_module(name + 'activation', activation) 112 | 113 | self.add_module(name + 'conv', conv_unit) 114 | 115 | if not preact: 116 | if bn: 117 | self.add_module(name + 'bn', bn_unit) 118 | 119 | if activation is not None: 120 | self.add_module(name + 'activation', activation) 121 | 122 | 123 | class Conv1d(_ConvBase): 124 | 125 | def __init__( 126 | self, 127 | in_size: int, 128 | out_size: int, 129 | *, 130 | kernel_size: int = 1, 131 | stride: int = 1, 132 | padding: int = 0, 133 | activation=nn.ReLU(inplace=True), 134 | bn: bool = False, 135 | init=nn.init.kaiming_normal_, 136 | bias: bool = True, 137 | preact: bool = False, 138 | name: str = "" 139 | ): 140 | super().__init__( 141 | in_size, 142 | out_size, 143 | kernel_size, 144 | stride, 145 | padding, 146 | activation, 147 | bn, 148 | init, 149 | conv=nn.Conv1d, 150 | batch_norm=BatchNorm1d, 151 | bias=bias, 152 | preact=preact, 153 | name=name 154 | ) 155 | 156 | 157 | class Conv2d(_ConvBase): 158 | 159 | def __init__( 160 | self, 161 | in_size: int, 162 | out_size: int, 163 | *, 164 | kernel_size: Tuple[int, int] = (1, 1), 165 | stride: Tuple[int, int] = (1, 1), 166 | padding: Tuple[int, int] = (0, 0), 167 | activation=nn.ReLU(inplace=True), 168 | bn: bool = False, 169 | init=nn.init.kaiming_normal_, 170 | bias: bool = True, 171 | preact: bool = False, 172 | name: str = "" 173 | ): 174 | super().__init__( 175 | in_size, 176 | out_size, 177 | kernel_size, 178 | stride, 179 | padding, 180 | activation, 181 | bn, 182 | init, 183 | conv=nn.Conv2d, 184 | batch_norm=BatchNorm2d, 185 | bias=bias, 186 | preact=preact, 187 | name=name 188 | ) 189 | 190 | 191 | class Conv3d(_ConvBase): 192 | 193 | def __init__( 194 | self, 195 | in_size: int, 196 | out_size: int, 197 | *, 198 | kernel_size: Tuple[int, int, int] = (1, 1, 1), 199 | stride: Tuple[int, int, int] = (1, 1, 1), 200 | padding: Tuple[int, int, int] = (0, 0, 0), 201 | activation=nn.ReLU(inplace=True), 202 | bn: bool = False, 203 | init=nn.init.kaiming_normal_, 204 | bias: bool = True, 205 | preact: bool = False, 206 | name: str = "" 207 | ): 208 | super().__init__( 209 | in_size, 210 | out_size, 211 | kernel_size, 212 | stride, 213 | padding, 214 | activation, 215 | bn, 216 | init, 217 | conv=nn.Conv3d, 218 | batch_norm=BatchNorm3d, 219 | bias=bias, 220 | preact=preact, 221 | name=name 222 | ) 223 | 224 | 225 | class FC(nn.Sequential): 226 | 227 | def __init__( 228 | self, 229 | in_size: int, 230 | out_size: int, 231 | *, 232 | activation=nn.ReLU(inplace=True), 233 | bn: bool = False, 234 | init=None, 235 | preact: bool = False, 236 | name: str = "" 237 | ): 238 | super().__init__() 239 | 240 | fc = nn.Linear(in_size, out_size, bias=not bn) 241 | if init is not None: 242 | init(fc.weight) 243 | if not bn: 244 | nn.init.constant_(fc.bias, 0) 245 | 246 | if preact: 247 | if bn: 248 | self.add_module(name + 'bn', BatchNorm1d(in_size)) 249 | 250 | if activation is not None: 251 | self.add_module(name + 'activation', activation) 252 | 253 | self.add_module(name + 'fc', fc) 254 | 255 | if not preact: 256 | if bn: 257 | self.add_module(name + 'bn', BatchNorm1d(out_size)) 258 | 259 | if activation is not None: 260 | self.add_module(name + 'activation', activation) 261 | 262 | def set_bn_momentum_default(bn_momentum): 263 | 264 | def fn(m): 265 | if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)): 266 | m.momentum = bn_momentum 267 | 268 | return fn 269 | 270 | 271 | class BNMomentumScheduler(object): 272 | 273 | def __init__( 274 | self, model, bn_lambda, last_epoch=-1, 275 | setter=set_bn_momentum_default 276 | ): 277 | if not isinstance(model, nn.Module): 278 | raise RuntimeError( 279 | "Class '{}' is not a PyTorch nn Module".format( 280 | type(model).__name__ 281 | ) 282 | ) 283 | 284 | self.model = model 285 | self.setter = setter 286 | self.lmbd = bn_lambda 287 | 288 | self.step(last_epoch + 1) 289 | self.last_epoch = last_epoch 290 | 291 | def step(self, epoch=None): 292 | if epoch is None: 293 | epoch = self.last_epoch + 1 294 | 295 | self.last_epoch = epoch 296 | self.model.apply(self.setter(self.lmbd(epoch))) 297 | 298 | 299 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | """ Testing for GraspNet baseline model. """ 2 | 3 | import os 4 | import sys 5 | import numpy as np 6 | import argparse 7 | import time 8 | from numba import jit 9 | from tqdm import tqdm 10 | import torch 11 | from torch.utils.data import DataLoader 12 | from graspnetAPI import GraspGroup, GraspNetEval 13 | 14 | ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) 15 | sys.path.append(os.path.join(ROOT_DIR, 'models')) 16 | sys.path.append(os.path.join(ROOT_DIR, 'dataset')) 17 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 18 | 19 | from models.FGC_graspnet import FGC_graspnet 20 | from models.loss import pred_decode 21 | from graspnet_dataset import GraspNetDataset, collate_fn, load_grasp_labels 22 | from collision_detector import ModelFreeCollisionDetector 23 | 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--dataset_root', default='./grasp_data', help='Dataset root') 26 | parser.add_argument('--checkpoint_path', default='/home/luyh/graspnet-baseline/logs_7155/best_noglobal/checkpoint.tar', help='Model checkpoint path') 27 | parser.add_argument('--camera', default='realsense', help='Camera split [realsense/kinect]') 28 | parser.add_argument('--num_point', type=int, default=20000, help='Point Number [default: 20000]') 29 | parser.add_argument('--num_view', type=int, default=300, help='View Number [default: 300]') 30 | parser.add_argument('--batch_size', type=int, default=2, help='Batch Size during inference [default: 1]') 31 | parser.add_argument('--collision_thresh', type=float, default=0.01, help='Collision Threshold in collision detection [default: 0.01]') 32 | parser.add_argument('--voxel_size', type=float, default=0.01, help='Voxel Size to process point clouds before collision detection [default: 0.01]') 33 | parser.add_argument('--num_workers', type=int, default=64, help='Number of workers used in evaluation [default: 30]') 34 | cfgs = parser.parse_args() 35 | 36 | 37 | # Init datasets and dataloaders 38 | def my_worker_init_fn(worker_id): 39 | np.random.seed(np.random.get_state()[1][0] + worker_id) 40 | pass 41 | 42 | # Create Dataset and Dataloader 43 | valid_obj_idxs, grasp_labels = load_grasp_labels(cfgs.dataset_root) 44 | #TEST_DATASET = GraspNetDataset(cfgs.dataset_root, valid_obj_idxs=None, grasp_labels=False, split='test', camera=cfgs.camera, num_points=cfgs.num_point, remove_outlier=True, augment=False, load_label=False) 45 | TEST_DATASET = GraspNetDataset(cfgs.dataset_root, valid_obj_idxs, grasp_labels, split='test', camera=cfgs.camera, num_points=cfgs.num_point, remove_outlier=True, augment=False, load_label=True) 46 | 47 | print(len(TEST_DATASET)) 48 | SCENE_LIST = TEST_DATASET.scene_list() 49 | TEST_DATALOADER = DataLoader(TEST_DATASET, batch_size=cfgs.batch_size, shuffle=False, 50 | num_workers=4, worker_init_fn=my_worker_init_fn, collate_fn=collate_fn) 51 | print(len(TEST_DATALOADER)) 52 | 53 | # Init the model 54 | net = FGC_graspnet(input_feature_dim=0, num_view=cfgs.num_view, num_angle=12, num_depth=4, 55 | cylinder_radius=0.05, hmin=-0.02, hmax=0.02, is_training=False, is_demo=False) 56 | 57 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 58 | net.to(device) 59 | 60 | # Load checkpoint 61 | checkpoint = torch.load(cfgs.checkpoint_path) 62 | net.load_state_dict(checkpoint['model_state_dict']) 63 | start_epoch = checkpoint['epoch'] 64 | print("-> loaded checkpoint %s (epoch: %d)"%(cfgs.checkpoint_path, start_epoch)) 65 | 66 | def ind_find(a, b): 67 | idx = np.where((a==b[:,None]).all(-1))[1] 68 | return idx 69 | 70 | def evaluate(): 71 | batch_interval = 256 72 | max_width = 0.1 73 | TOP_K = 50 74 | #score_active = [0.2, 0.4, 0.6, 0.8] 75 | score_active = [0, 0.1, 0.3, 0.5, 0.7, 0.9] 76 | stat_dict = {} # collect statistics 77 | # set model to eval mode (for bn and dp) 78 | net.eval() 79 | tic = time.time() 80 | from ipdb import set_trace 81 | grasp_accuracy_scene = [] 82 | for batch_idx, batch_data in enumerate(tqdm(TEST_DATALOADER)): 83 | for key in batch_data: 84 | if 'list' in key: 85 | for i in range(len(batch_data[key])): 86 | for j in range(len(batch_data[key][i])): 87 | batch_data[key][i][j] = batch_data[key][i][j].to(device) 88 | else: 89 | batch_data[key] = batch_data[key].to(device) 90 | 91 | # Forward pass 92 | with torch.no_grad(): 93 | end_points = net(batch_data) 94 | grasp_preds, score_labels = pred_decode(end_points) 95 | 96 | for i in range(cfgs.batch_size): 97 | data_idx = batch_idx * cfgs.batch_size + i 98 | preds = grasp_preds[i].detach().cpu().numpy() 99 | score_labels_ = score_labels[i].detach().cpu().numpy() 100 | gg = GraspGroup(preds) 101 | gg = gg.nms(0.03, 30.0 / 180 * np.pi) 102 | gg1_ = np.array(gg.grasp_group_array) 103 | ind1_ = ind_find(preds, gg1_) 104 | score_labels_ = score_labels_[ind1_] 105 | 106 | # collision detection and nms 107 | if cfgs.collision_thresh > 0: 108 | cloud, _ = TEST_DATASET.get_data(data_idx, return_raw_cloud=True) 109 | mfcdetector = ModelFreeCollisionDetector(cloud, voxel_size=cfgs.voxel_size) 110 | collision_mask = mfcdetector.detect(gg, approach_dist=0.05, collision_thresh=cfgs.collision_thresh) 111 | gg = gg[~collision_mask] 112 | score_labels_ = score_labels_[~collision_mask] 113 | preds_ = np.array(gg.grasp_group_array) 114 | 115 | 116 | #set_trace() 117 | # check the width of grasp prediction 118 | min_width_mask = (preds_[:, 1] < 0) 119 | max_width_mask = (preds_[:, 1] > max_width) 120 | preds_[min_width_mask, 1] = 0 121 | preds_[max_width_mask, 1] = max_width 122 | pred_score = preds_[:, 0] 123 | 124 | # sort the pred score 125 | idx = np.argsort(-pred_score) 126 | score_test = score_labels_[idx] 127 | 128 | grasp_accuracy = np.zeros((TOP_K, len(score_active))) 129 | for i, score_gap in enumerate(score_active): 130 | for k in range(TOP_K): 131 | if k+1> len(score_test): 132 | grasp_accuracy[k, i] = np.sum((score_test >=score_gap).astype(int))/(k+1) 133 | else: 134 | grasp_accuracy[k, i] = np.sum((score_test[0:k+1]>=score_gap).astype(int))/(k+1) 135 | grasp_accuracy_scene.append(grasp_accuracy) 136 | 137 | if (batch_idx+1) % 256 == 0: 138 | toc = time.time() 139 | print('Eval batch scene: %d/90, time: %fs' % (batch_idx//256, (toc - tic) / batch_interval)) 140 | tic = time.time() 141 | 142 | grasp_accuracy_scene = np.array(grasp_accuracy_scene) 143 | grasp_ap = np.reshape(grasp_accuracy_scene, (-1, 256, TOP_K, len(score_active))) 144 | AP = [100*np.mean(np.mean(np.mean(grasp_ap, axis=2), axis=1), axis=0), 145 | 100*np.mean(np.mean(np.mean(grasp_ap[0:30], axis=2), axis=1), axis=0), 146 | 100*np.mean(np.mean(np.mean(grasp_ap[30:60], axis=2), axis=1), axis=0), 147 | 100*np.mean(np.mean(np.mean(grasp_ap[60:90], axis=2), axis=1), axis=0)] 148 | mAP = np.mean(np.array(AP), axis=1) 149 | print('\nEvaluation Result:\n----------\n{}, mAP={}, mAP Seen={}, mAP Similar={}, mAP Novel={}'.format('realsense', mAP[0], mAP[1], mAP[2], mAP[3])) 150 | print('\nEvaluation Result:\n----------\n{}, AP0.7={}, AP0.7 Seen={}, AP0.7 Similar={}, AP0.7 Novel={}'.format('realsense', 151 | AP[0][-2], 152 | AP[1][-2], 153 | AP[2][-2], 154 | AP[3][-2])) 155 | print('\nEvaluation Result:\n----------\n{}, AP0.3={}, AP0.3 Seen={}, AP0.3 Similar={}, AP0.3 Novel={}'.format('realsense', 156 | AP[0][2], 157 | AP[1][2], 158 | AP[2][2], 159 | AP[3][2])) 160 | if __name__=='__main__': 161 | evaluate() 162 | -------------------------------------------------------------------------------- /models/modules.py: -------------------------------------------------------------------------------- 1 | """ Modules for GraspNet baseline model. 2 | Author: chenxi-wang 3 | """ 4 | 5 | import os 6 | import sys 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as F 10 | 11 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 12 | ROOT_DIR = os.path.dirname(BASE_DIR) 13 | sys.path.append(ROOT_DIR) 14 | sys.path.append(os.path.join(ROOT_DIR, 'pointnet2')) 15 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 16 | 17 | import pytorch_utils as pt_utils 18 | from pointnet2.pointnet2_utils import CylinderQueryAndGroup 19 | from loss_utils import generate_grasp_views, batch_viewpoint_params_to_matrix 20 | 21 | 22 | class ApproachNet(nn.Module): 23 | def __init__(self, num_view, seed_feature_dim): 24 | """ Approach vector estimation from seed point features. 25 | 26 | Input: 27 | num_view: [int] 28 | number of views generated from each each seed point 29 | seed_feature_dim: [int] 30 | number of channels of seed point features 31 | """ 32 | super().__init__() 33 | self.num_view = num_view # 300 34 | self.in_dim = seed_feature_dim # 256 35 | self.conv1 = nn.Conv1d(self.in_dim, self.in_dim, 1) 36 | self.conv2 = nn.Conv1d(self.in_dim, 2+self.num_view, 1) 37 | self.conv3 = nn.Conv1d(2+self.num_view, 2+self.num_view, 1) 38 | self.bn1 = nn.BatchNorm1d(self.in_dim) 39 | self.bn2 = nn.BatchNorm1d(2+self.num_view) 40 | 41 | def forward(self, seed_xyz, seed_features, end_points): 42 | """ Forward pass. 43 | 44 | Input: 45 | seed_xyz: [torch.FloatTensor, (batch_size,num_seed,3)] 46 | coordinates of seed points 47 | seed_features: [torch.FloatTensor, (batch_size,feature_dim,num_seed) 48 | features of seed points 49 | end_points: [dict] 50 | 51 | Output: 52 | end_points: [dict] 53 | """ 54 | B, num_seed, _ = seed_xyz.size() 55 | features = F.relu(self.bn1(self.conv1(seed_features)), inplace=True) 56 | features = F.relu(self.bn2(self.conv2(features)), inplace=True) 57 | features = self.conv3(features) 58 | objectness_score = features[:, :2, :] # (B, 2, num_seed) 59 | view_score = features[:, 2:2+self.num_view, :].transpose(1,2).contiguous() # (B, num_seed, num_view) 60 | end_points['objectness_score'] = objectness_score 61 | end_points['view_score'] = view_score 62 | 63 | # print(view_score.min(), view_score.max(), view_score.mean()) 64 | top_view_scores, top_view_inds = torch.max(view_score, dim=2) # (B, num_seed) 65 | top_view_inds_ = top_view_inds.view(B, num_seed, 1, 1).expand(-1, -1, -1, 3).contiguous() 66 | template_views = generate_grasp_views(self.num_view).to(features.device) # (num_view, 3) 67 | template_views = template_views.view(1, 1, self.num_view, 3).expand(B, num_seed, -1, -1).contiguous() #(B, num_seed, num_view, 3) 68 | vp_xyz = torch.gather(template_views, 2, top_view_inds_).squeeze(2) #(B, num_seed, 3) 69 | vp_xyz_ = vp_xyz.view(-1, 3) 70 | batch_angle = torch.zeros(vp_xyz_.size(0), dtype=vp_xyz.dtype, device=vp_xyz.device) 71 | vp_rot = batch_viewpoint_params_to_matrix(-vp_xyz_, batch_angle).view(B, num_seed, 3, 3) 72 | end_points['grasp_top_view_inds'] = top_view_inds 73 | end_points['grasp_top_view_score'] = top_view_scores 74 | end_points['grasp_top_view_xyz'] = vp_xyz 75 | end_points['grasp_top_view_rot'] = vp_rot 76 | 77 | return end_points 78 | 79 | 80 | class CloudCrop(nn.Module): 81 | """ Cylinder group and align for grasp configure estimation. Return a list of grouped points with different cropping depths. 82 | 83 | Input: 84 | nsample: [int] 85 | sample number in a group 86 | seed_feature_dim: [int] 87 | number of channels of grouped points 88 | cylinder_radius: [float] 89 | radius of the cylinder space 90 | hmin: [float] 91 | height of the bottom surface 92 | hmax_list: [list of float] 93 | list of heights of the upper surface 94 | """ 95 | def __init__(self, nsample, seed_feature_dim, cylinder_radius=0.05, hmin=-0.02, hmax_list=[0.01,0.02,0.03,0.04]): 96 | super().__init__() 97 | self.nsample = nsample # 64 98 | self.in_dim = seed_feature_dim # 3 99 | self.cylinder_radius = cylinder_radius 100 | mlps = [self.in_dim, 64, 128, 256] 101 | 102 | self.groupers = [] 103 | for hmax in hmax_list: 104 | self.groupers.append(CylinderQueryAndGroup( 105 | cylinder_radius, hmin, hmax, nsample, use_xyz=True 106 | )) 107 | self.mlps = pt_utils.SharedMLP(mlps, bn=True) 108 | #self.local_att = Local_attention(256) 109 | 110 | 111 | def forward(self, seed_xyz, pointcloud, vp_rot): 112 | """ Forward pass. 113 | 114 | Input: 115 | seed_xyz: [torch.FloatTensor, (batch_size,num_seed,3)] 116 | coordinates of seed points 117 | pointcloud: [torch.FloatTensor, (batch_size,num_seed,3)] 118 | the points to be cropped 119 | vp_rot: [torch.FloatTensor, (batch_size,num_seed,3,3)] 120 | rotation matrices generated from approach vectors 121 | 122 | Output: 123 | vp_features: [torch.FloatTensor, (batch_size,num_features,num_seed,num_depth)] 124 | features of grouped points in different depths 125 | """ 126 | B, num_seed, _, _ = vp_rot.size() 127 | num_depth = len(self.groupers) 128 | grouped_features = [] 129 | for grouper in self.groupers: 130 | grouped_features.append(grouper( 131 | pointcloud, seed_xyz, vp_rot 132 | )) # (batch_size, feature_dim, num_seed, nsample) 133 | grouped_features = torch.stack(grouped_features, dim=3) # (batch_size, feature_dim, num_seed, num_depth, nsample) 134 | grouped_features = grouped_features.view(B, -1, num_seed*num_depth, self.nsample) # (batch_size, feature_dim, num_seed*num_depth, nsample) 135 | 136 | 137 | vp_features = self.mlps( 138 | grouped_features 139 | ) # (batch_size, mlps[-1], num_seed*num_depth, nsample) 140 | 141 | # vp_features = vp_features.permute(0, 2, 1, 3).contiguous().view(B*num_seed*num_depth, 256, self.nsample) # (B*num_seed*num_depth, C, K) 142 | # vp_features = self.local_att(vp_features).contiguous().view(B, num_seed*num_depth, 256, self.nsample).permute(0, 2, 1, 3) 143 | 144 | vp_features = F.max_pool2d( 145 | vp_features, kernel_size=[1, vp_features.size(3)] 146 | ) # (batch_size, mlps[-1], num_seed*num_depth, 1) 147 | vp_features = vp_features.view(B, -1, num_seed, num_depth) 148 | return vp_features 149 | 150 | 151 | class OperationNet(nn.Module): 152 | """ Grasp configure estimation. 153 | Input: 154 | num_angle: [int] 155 | number of in-plane rotation angle classes 156 | the value of the i-th class --> i*PI/num_angle (i=0,...,num_angle-1) 157 | num_depth: [int] 158 | number of gripper depth classes 159 | """ 160 | def __init__(self, num_angle, num_depth): 161 | # Output: 162 | # scores(num_angle) 163 | # angle class (num_angle) 164 | # width (num_angle) 165 | super().__init__() 166 | self.num_angle = num_angle 167 | self.num_depth = num_depth 168 | 169 | self.conv1 = nn.Conv1d(256, 128, 1) 170 | self.conv2 = nn.Conv1d(128, 128, 1) 171 | self.conv3 = nn.Conv1d(128, 3*num_angle, 1) 172 | self.bn1 = nn.BatchNorm1d(128) 173 | self.bn2 = nn.BatchNorm1d(128) 174 | 175 | def forward(self, vp_features, end_points): 176 | """ Forward pass. 177 | Input: 178 | vp_features: [torch.FloatTensor, (batch_size,num_seed,3)] 179 | features of grouped points in different depths 180 | end_points: [dict] 181 | Output: 182 | end_points: [dict] 183 | """ 184 | B, _, num_seed, num_depth = vp_features.size() 185 | vp_features = vp_features.view(B, -1, num_seed*num_depth) 186 | vp_features = F.relu(self.bn1(self.conv1(vp_features)), inplace=True) 187 | vp_features = F.relu(self.bn2(self.conv2(vp_features)), inplace=True) 188 | vp_features = self.conv3(vp_features) 189 | vp_features = vp_features.view(B, -1, num_seed, num_depth) 190 | 191 | # split prediction 192 | end_points['grasp_score_pred'] = vp_features[:, 0:self.num_angle] 193 | end_points['grasp_angle_cls_pred'] = vp_features[:, self.num_angle:2*self.num_angle] 194 | end_points['grasp_width_pred'] = vp_features[:, 2*self.num_angle:3*self.num_angle] 195 | return end_points -------------------------------------------------------------------------------- /knn/src/cuda/knn.cu: -------------------------------------------------------------------------------- 1 | /** Modifed version of knn-CUDA from https://github.com/vincentfpgarcia/kNN-CUDA 2 | * The modifications are 3 | * removed texture memory usage 4 | * removed split query KNN computation 5 | * added feature extraction with bilinear interpolation 6 | * 7 | * Last modified by Christopher B. Choy 12/23/2016 8 | */ 9 | 10 | // Includes 11 | #include 12 | #include "cuda.h" 13 | 14 | #define IDX2D(i, j, dj) (dj * i + j) 15 | #define IDX3D(i, j, k, dj, dk) (IDX2D(IDX2D(i, j, dj), k, dk)) 16 | 17 | #define BLOCK 512 18 | #define MAX_STREAMS 512 19 | 20 | // Constants used by the program 21 | #define BLOCK_DIM 16 22 | #define DEBUG 0 23 | 24 | 25 | /** 26 | * Computes the distance between two matrix A (reference points) and 27 | * B (query points) containing respectively wA and wB points. 28 | * 29 | * @param A pointer on the matrix A 30 | * @param wA width of the matrix A = number of points in A 31 | * @param B pointer on the matrix B 32 | * @param wB width of the matrix B = number of points in B 33 | * @param dim dimension of points = height of matrices A and B 34 | * @param AB pointer on the matrix containing the wA*wB distances computed 35 | */ 36 | __global__ void cuComputeDistanceGlobal( float* A, int wA, 37 | float* B, int wB, int dim, float* AB){ 38 | 39 | // Declaration of the shared memory arrays As and Bs used to store the sub-matrix of A and B 40 | __shared__ float shared_A[BLOCK_DIM][BLOCK_DIM]; 41 | __shared__ float shared_B[BLOCK_DIM][BLOCK_DIM]; 42 | 43 | 44 | // Sub-matrix of A (begin, step, end) and Sub-matrix of B (begin, step) 45 | __shared__ int begin_A; 46 | __shared__ int begin_B; 47 | __shared__ int step_A; 48 | __shared__ int step_B; 49 | __shared__ int end_A; 50 | 51 | // Thread index 52 | int tx = threadIdx.x; 53 | int ty = threadIdx.y; 54 | 55 | // Other variables 56 | float tmp; 57 | float ssd = 0; 58 | 59 | // Loop parameters 60 | begin_A = BLOCK_DIM * blockIdx.y; 61 | begin_B = BLOCK_DIM * blockIdx.x; 62 | step_A = BLOCK_DIM * wA; 63 | step_B = BLOCK_DIM * wB; 64 | end_A = begin_A + (dim-1) * wA; 65 | 66 | // Conditions 67 | int cond0 = (begin_A + tx < wA); // used to write in shared memory 68 | int cond1 = (begin_B + tx < wB); // used to write in shared memory & to computations and to write in output matrix 69 | int cond2 = (begin_A + ty < wA); // used to computations and to write in output matrix 70 | 71 | // Loop over all the sub-matrices of A and B required to compute the block sub-matrix 72 | for (int a = begin_A, b = begin_B; a <= end_A; a += step_A, b += step_B) { 73 | // Load the matrices from device memory to shared memory; each thread loads one element of each matrix 74 | if (a/wA + ty < dim){ 75 | shared_A[ty][tx] = (cond0)? A[a + wA * ty + tx] : 0; 76 | shared_B[ty][tx] = (cond1)? B[b + wB * ty + tx] : 0; 77 | } 78 | else{ 79 | shared_A[ty][tx] = 0; 80 | shared_B[ty][tx] = 0; 81 | } 82 | 83 | // Synchronize to make sure the matrices are loaded 84 | __syncthreads(); 85 | 86 | // Compute the difference between the two matrixes; each thread computes one element of the block sub-matrix 87 | if (cond2 && cond1){ 88 | for (int k = 0; k < BLOCK_DIM; ++k){ 89 | tmp = shared_A[k][ty] - shared_B[k][tx]; 90 | ssd += tmp*tmp; 91 | } 92 | } 93 | 94 | // Synchronize to make sure that the preceding computation is done before loading two new sub-matrices of A and B in the next iteration 95 | __syncthreads(); 96 | } 97 | 98 | // Write the block sub-matrix to device memory; each thread writes one element 99 | if (cond2 && cond1) 100 | AB[(begin_A + ty) * wB + begin_B + tx] = ssd; 101 | } 102 | 103 | 104 | /** 105 | * Gathers k-th smallest distances for each column of the distance matrix in the top. 106 | * 107 | * @param dist distance matrix 108 | * @param ind index matrix 109 | * @param width width of the distance matrix and of the index matrix 110 | * @param height height of the distance matrix and of the index matrix 111 | * @param k number of neighbors to consider 112 | */ 113 | __global__ void cuInsertionSort(float *dist, long *ind, int width, int height, int k){ 114 | 115 | // Variables 116 | int l, i, j; 117 | float *p_dist; 118 | long *p_ind; 119 | float curr_dist, max_dist; 120 | long curr_row, max_row; 121 | unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x; 122 | 123 | if (xIndexcurr_dist){ 138 | i=a; 139 | break; 140 | } 141 | } 142 | for (j=l; j>i; j--){ 143 | p_dist[j*width] = p_dist[(j-1)*width]; 144 | p_ind[j*width] = p_ind[(j-1)*width]; 145 | } 146 | p_dist[i*width] = curr_dist; 147 | p_ind[i*width] = l+1; 148 | } else { 149 | p_ind[l*width] = l+1; 150 | } 151 | max_dist = p_dist[curr_row]; 152 | } 153 | 154 | // Part 2 : insert element in the k-th first lines 155 | max_row = (k-1)*width; 156 | for (l=k; lcurr_dist){ 162 | i=a; 163 | break; 164 | } 165 | } 166 | for (j=k-1; j>i; j--){ 167 | p_dist[j*width] = p_dist[(j-1)*width]; 168 | p_ind[j*width] = p_ind[(j-1)*width]; 169 | } 170 | p_dist[i*width] = curr_dist; 171 | p_ind[i*width] = l+1; 172 | max_dist = p_dist[max_row]; 173 | } 174 | } 175 | } 176 | } 177 | 178 | 179 | /** 180 | * Computes the square root of the first line (width-th first element) 181 | * of the distance matrix. 182 | * 183 | * @param dist distance matrix 184 | * @param width width of the distance matrix 185 | * @param k number of neighbors to consider 186 | */ 187 | __global__ void cuParallelSqrt(float *dist, int width, int k){ 188 | unsigned int xIndex = blockIdx.x * blockDim.x + threadIdx.x; 189 | unsigned int yIndex = blockIdx.y * blockDim.y + threadIdx.y; 190 | if (xIndex>>(ref_dev, ref_nb, query_dev, query_nb, dim, dist_dev); 237 | 238 | // Kernel 2: Sort each column 239 | cuInsertionSort<<>>(dist_dev, ind_dev, query_nb, ref_nb, k); 240 | 241 | // Kernel 3: Compute square root of k first elements 242 | // cuParallelSqrt<<>>(dist_dev, query_nb, k); 243 | 244 | #if DEBUG 245 | unsigned int size_of_float = sizeof(float); 246 | unsigned long size_of_long = sizeof(long); 247 | 248 | float* dist_host = new float[query_nb * k]; 249 | long* idx_host = new long[query_nb * k]; 250 | 251 | // Memory copy of output from device to host 252 | cudaMemcpy(&dist_host[0], dist_dev, 253 | query_nb * k *size_of_float, cudaMemcpyDeviceToHost); 254 | 255 | cudaMemcpy(&idx_host[0], ind_dev, 256 | query_nb * k * size_of_long, cudaMemcpyDeviceToHost); 257 | 258 | int i = 0; 259 | for(i = 0; i < 100; i++){ 260 | printf("IDX[%d]: %d\n", i, (int)idx_host[i]); 261 | } 262 | #endif 263 | } 264 | 265 | 266 | 267 | 268 | 269 | 270 | -------------------------------------------------------------------------------- /log/log_train.txt: -------------------------------------------------------------------------------- 1 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 2 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 3 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 4 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 5 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='~/graspnet_baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 6 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='~/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 7 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 8 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 9 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 10 | **** EPOCH 000 **** 11 | Current learning rate: 0.001000 12 | Current BN decay momentum: 0.500000 13 | 2022-07-06 22:53:02.561044 14 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 15 | **** EPOCH 000 **** 16 | Current learning rate: 0.001000 17 | Current BN decay momentum: 0.500000 18 | 2022-07-06 22:54:20.664333 19 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 20 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 21 | **** EPOCH 000 **** 22 | Current learning rate: 0.001000 23 | Current BN decay momentum: 0.500000 24 | 2022-07-06 22:58:01.802845 25 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 26 | **** EPOCH 000 **** 27 | Current learning rate: 0.001000 28 | Current BN decay momentum: 0.500000 29 | 2022-07-06 23:00:48.033211 30 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 31 | **** EPOCH 000 **** 32 | Current learning rate: 0.001000 33 | Current BN decay momentum: 0.500000 34 | 2022-07-06 23:02:30.260290 35 | ---- batch: 010 ---- 36 | mean loss/overall_loss: 1.777084 37 | mean loss/stage1_objectness_loss: 0.650260 38 | mean loss/stage1_view_loss: 0.141157 39 | mean loss/stage2_grasp_angle_class_loss: 2.166051 40 | mean loss/stage2_grasp_depth_cls_loss: 0.968530 41 | mean loss/stage2_grasp_score_loss: 0.116216 42 | mean loss/stage2_grasp_width_loss: 1.329874 43 | mean stage1_objectness_acc: 0.626758 44 | mean stage1_objectness_prec: 0.576089 45 | mean stage1_objectness_recall: 0.578866 46 | mean stage1_pos_view_pred_count: 6743.100000 47 | mean stage2_grasp_angle_class_acc/0_degree: 0.421704 48 | mean stage2_grasp_angle_class_acc/15_degree: 0.522584 49 | mean stage2_grasp_angle_class_acc/30_degree: 0.620284 50 | mean stage2_grasp_depth_cls_acc: 0.733744 51 | ---- batch: 020 ---- 52 | mean loss/overall_loss: 1.280322 53 | mean loss/stage1_objectness_loss: 0.560989 54 | mean loss/stage1_view_loss: 0.083024 55 | mean loss/stage2_grasp_angle_class_loss: 1.502694 56 | mean loss/stage2_grasp_depth_cls_loss: 0.660265 57 | mean loss/stage2_grasp_score_loss: 0.057837 58 | mean loss/stage2_grasp_width_loss: 0.652509 59 | mean stage1_objectness_acc: 0.699658 60 | mean stage1_objectness_prec: 0.602643 61 | mean stage1_objectness_recall: 0.649153 62 | mean stage1_pos_view_pred_count: 1151.000000 63 | mean stage2_grasp_angle_class_acc/0_degree: 0.770204 64 | mean stage2_grasp_angle_class_acc/15_degree: 0.808371 65 | mean stage2_grasp_angle_class_acc/30_degree: 0.839097 66 | mean stage2_grasp_depth_cls_acc: 0.834418 67 | ---- batch: 030 ---- 68 | mean loss/overall_loss: 1.092311 69 | mean loss/stage1_objectness_loss: 0.580696 70 | mean loss/stage1_view_loss: 0.050878 71 | mean loss/stage2_grasp_angle_class_loss: 1.147558 72 | mean loss/stage2_grasp_depth_cls_loss: 0.522689 73 | mean loss/stage2_grasp_score_loss: 0.040676 74 | mean loss/stage2_grasp_width_loss: 0.419882 75 | mean stage1_objectness_acc: 0.702441 76 | mean stage1_objectness_prec: 0.664283 77 | mean stage1_objectness_recall: 0.722532 78 | mean stage1_pos_view_pred_count: 328.800000 79 | mean stage2_grasp_angle_class_acc/0_degree: 0.826170 80 | mean stage2_grasp_angle_class_acc/15_degree: 0.852629 81 | mean stage2_grasp_angle_class_acc/30_degree: 0.881014 82 | mean stage2_grasp_depth_cls_acc: 0.833918 83 | ---- batch: 040 ---- 84 | mean loss/overall_loss: 0.945058 85 | mean loss/stage1_objectness_loss: 0.541682 86 | mean loss/stage1_view_loss: 0.039571 87 | mean loss/stage2_grasp_angle_class_loss: 0.893422 88 | mean loss/stage2_grasp_depth_cls_loss: 0.441928 89 | mean loss/stage2_grasp_score_loss: 0.026413 90 | mean loss/stage2_grasp_width_loss: 0.307784 91 | mean stage1_objectness_acc: 0.721338 92 | mean stage1_objectness_prec: 0.668028 93 | mean stage1_objectness_recall: 0.735213 94 | mean stage1_pos_view_pred_count: 105.600000 95 | mean stage2_grasp_angle_class_acc/0_degree: 0.842629 96 | mean stage2_grasp_angle_class_acc/15_degree: 0.868705 97 | mean stage2_grasp_angle_class_acc/30_degree: 0.889179 98 | mean stage2_grasp_depth_cls_acc: 0.842629 99 | ---- batch: 050 ---- 100 | mean loss/overall_loss: 0.986783 101 | mean loss/stage1_objectness_loss: 0.542187 102 | mean loss/stage1_view_loss: 0.040517 103 | mean loss/stage2_grasp_angle_class_loss: 0.993633 104 | mean loss/stage2_grasp_depth_cls_loss: 0.548309 105 | mean loss/stage2_grasp_score_loss: 0.022733 106 | mean loss/stage2_grasp_width_loss: 0.252751 107 | mean stage1_objectness_acc: 0.710693 108 | mean stage1_objectness_prec: 0.608728 109 | mean stage1_objectness_recall: 0.771846 110 | mean stage1_pos_view_pred_count: 46.900000 111 | mean stage2_grasp_angle_class_acc/0_degree: 0.786961 112 | mean stage2_grasp_angle_class_acc/15_degree: 0.824893 113 | mean stage2_grasp_angle_class_acc/30_degree: 0.851668 114 | mean stage2_grasp_depth_cls_acc: 0.786961 115 | ---- batch: 060 ---- 116 | mean loss/overall_loss: 0.990469 117 | mean loss/stage1_objectness_loss: 0.568085 118 | mean loss/stage1_view_loss: 0.036698 119 | mean loss/stage2_grasp_angle_class_loss: 0.967829 120 | mean loss/stage2_grasp_depth_cls_loss: 0.531584 121 | mean loss/stage2_grasp_score_loss: 0.019485 122 | mean loss/stage2_grasp_width_loss: 0.212771 123 | mean stage1_objectness_acc: 0.701611 124 | mean stage1_objectness_prec: 0.645445 125 | mean stage1_objectness_recall: 0.673493 126 | mean stage1_pos_view_pred_count: 17.700000 127 | mean stage2_grasp_angle_class_acc/0_degree: 0.790531 128 | mean stage2_grasp_angle_class_acc/15_degree: 0.827681 129 | mean stage2_grasp_angle_class_acc/30_degree: 0.857589 130 | mean stage2_grasp_depth_cls_acc: 0.790531 131 | Namespace(batch_size=2, bn_decay_rate=0.5, bn_decay_step=2, camera='realsense', checkpoint_path=None, dataset_root='/home/luyh/graspnet-baseline/grasp_data', device='cuda', dist_url='env://', distributed=False, learning_rate=0.001, log_dir='log', lr_decay_rates='0.1,0.1,0.1', lr_decay_steps='8,12,16', max_epoch=18, num_point=20000, num_view=300, seed=13, weight_decay=0, world_size=1) 132 | **** EPOCH 000 **** 133 | Current learning rate: 0.001000 134 | Current BN decay momentum: 0.500000 135 | 2022-07-06 23:14:34.317510 136 | ---- batch: 010 ---- 137 | mean loss/overall_loss: 1.772583 138 | mean loss/stage1_objectness_loss: 0.644728 139 | mean loss/stage1_view_loss: 0.140702 140 | mean loss/stage2_grasp_angle_class_loss: 2.173053 141 | mean loss/stage2_grasp_depth_cls_loss: 0.974891 142 | mean loss/stage2_grasp_score_loss: 0.116768 143 | mean loss/stage2_grasp_width_loss: 1.320168 144 | mean stage1_objectness_acc: 0.630518 145 | mean stage1_objectness_prec: 0.579949 146 | mean stage1_objectness_recall: 0.584234 147 | mean stage1_pos_view_pred_count: 6654.700000 148 | mean stage2_grasp_angle_class_acc/0_degree: 0.413733 149 | mean stage2_grasp_angle_class_acc/15_degree: 0.517574 150 | mean stage2_grasp_angle_class_acc/30_degree: 0.617114 151 | mean stage2_grasp_depth_cls_acc: 0.729804 152 | -------------------------------------------------------------------------------- /models/loss.py: -------------------------------------------------------------------------------- 1 | """ Loss functions for training. 2 | Author: chenxi-wang 3 | """ 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import numpy as np 9 | import sys 10 | import os 11 | import time 12 | 13 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 14 | ROOT_DIR = os.path.dirname(BASE_DIR) 15 | sys.path.append(ROOT_DIR) 16 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 17 | 18 | from utils.loss_utils import GRASP_MAX_WIDTH, GRASP_MAX_TOLERANCE, THRESH_GOOD, THRESH_BAD,\ 19 | transform_point_cloud, generate_grasp_views,\ 20 | batch_viewpoint_params_to_matrix, huber_loss 21 | 22 | def get_loss(end_points): 23 | objectness_loss, end_points = compute_objectness_loss(end_points) 24 | view_loss, end_points = compute_view_loss(end_points) 25 | grasp_loss, end_points = compute_grasp_loss(end_points) 26 | loss = objectness_loss + 3*view_loss + 0.2*grasp_loss 27 | end_points['loss/overall_loss'] = loss 28 | return loss, end_points 29 | 30 | def compute_objectness_loss(end_points): 31 | criterion = nn.CrossEntropyLoss(reduction='mean') 32 | objectness_score = end_points['objectness_score'] 33 | objectness_label = end_points['objectness_label'] 34 | fp2_inds = end_points['fp2_inds'].long() 35 | objectness_label = torch.gather(objectness_label, 1, fp2_inds) 36 | loss = criterion(objectness_score, objectness_label) 37 | 38 | end_points['loss/stage1_objectness_loss'] = loss 39 | objectness_pred = torch.argmax(objectness_score, 1) 40 | end_points['stage1_objectness_acc'] = (objectness_pred == objectness_label.long()).float().mean() 41 | 42 | end_points['stage1_objectness_prec'] = (objectness_pred == objectness_label.long())[objectness_pred == 1].float().mean() 43 | end_points['stage1_objectness_recall'] = (objectness_pred == objectness_label.long())[objectness_label == 1].float().mean() 44 | 45 | return loss, end_points 46 | 47 | def compute_view_loss(end_points): 48 | criterion = nn.MSELoss(reduction='none') 49 | view_score = end_points['view_score'] 50 | view_label = end_points['batch_grasp_view_label'] 51 | objectness_label = end_points['objectness_label'] 52 | fp2_inds = end_points['fp2_inds'].long() 53 | V = view_label.size(2) 54 | objectness_label = torch.gather(objectness_label, 1, fp2_inds) 55 | 56 | objectness_mask = (objectness_label > 0) 57 | objectness_mask = objectness_mask.unsqueeze(-1).repeat(1, 1, V) 58 | pos_view_pred_mask = ((view_score >= THRESH_GOOD) & objectness_mask) 59 | 60 | loss = criterion(view_score, view_label) 61 | loss = loss[objectness_mask].mean() 62 | 63 | end_points['loss/stage1_view_loss'] = loss 64 | end_points['stage1_pos_view_pred_count'] = pos_view_pred_mask.long().sum() 65 | 66 | return loss, end_points 67 | 68 | 69 | def compute_grasp_loss(end_points, use_template_in_training=True): 70 | top_view_inds = end_points['grasp_top_view_inds'] # (B, Ns) 71 | vp_rot = end_points['grasp_top_view_rot'] # (B, Ns, view_factor, 3, 3) 72 | objectness_label = end_points['objectness_label'] 73 | fp2_inds = end_points['fp2_inds'].long() 74 | objectness_mask = torch.gather(objectness_label, 1, fp2_inds).bool() # (B, Ns) 75 | 76 | # process labels 77 | batch_grasp_label = end_points['batch_grasp_label'] # (B, Ns, A, D) 78 | batch_grasp_offset = end_points['batch_grasp_offset'] # (B, Ns, A, D, 3) 79 | B, Ns, A, D = batch_grasp_label.size() 80 | 81 | # pick the one with the highest angle score 82 | top_view_grasp_angles = batch_grasp_offset[:, :, :, :, 0] #(B, Ns, A, D) 83 | top_view_grasp_depths = batch_grasp_offset[:, :, :, :, 1] #(B, Ns, A, D) 84 | top_view_grasp_widths = batch_grasp_offset[:, :, :, :, 2] #(B, Ns, A, D) 85 | target_labels_inds = torch.argmax(batch_grasp_label, dim=2, keepdim=True) # (B, Ns, 1, D) 86 | #target_labels = torch.gather(batch_grasp_label, 2, target_labels_inds).squeeze(2) # (B, Ns, D) 87 | #target_angles = torch.gather(top_view_grasp_angles, 2, target_labels_inds).squeeze(2) # (B, Ns, D) 88 | #target_depths = torch.gather(top_view_grasp_depths, 2, target_labels_inds).squeeze(2) # (B, Ns, D) 89 | target_widths = torch.gather(top_view_grasp_widths, 2, target_labels_inds).squeeze(2) # (B, Ns, D) 90 | 91 | # graspable_mask = (target_labels > THRESH_BAD) 92 | # objectness_mask = objectness_mask.unsqueeze(-1).expand_as(graspable_mask) 93 | # loss_mask = (objectness_mask & graspable_mask).float() 94 | 95 | # 1. grasp score loss 96 | criterion_score = nn.MSELoss(reduction='none') 97 | target_scores = batch_grasp_label.view(B, Ns, -1) # B, N, 48 98 | grasp_score = end_points['grasp_score_pred'].transpose(1, 2).contiguous() # B, N, 48 99 | grasp_score_loss = criterion_score(grasp_score, target_scores) 100 | objectness_mask1 = objectness_mask.unsqueeze(-1).expand_as(target_scores) 101 | grasp_score_loss = grasp_score_loss[objectness_mask1].mean() 102 | end_points['loss/stage2_grasp_score_loss'] = grasp_score_loss 103 | 104 | # 2. inplane rotation cls loss 105 | target_angles_cls = torch.argmax(torch.max(batch_grasp_label, -1).values, -1) # B, N 106 | criterion_grasp_angle_class = nn.CrossEntropyLoss(reduction='none') 107 | grasp_angle_class_score = end_points['grasp_angle_cls_pred'] # B, 12, N 108 | grasp_angle_class_loss = criterion_grasp_angle_class(grasp_angle_class_score, target_angles_cls) 109 | 110 | grasp_angle_class_loss = grasp_angle_class_loss[objectness_mask].mean() 111 | end_points['loss/stage2_grasp_angle_class_loss'] = grasp_angle_class_loss 112 | grasp_angle_class_pred = torch.argmax(grasp_angle_class_score, 1) 113 | end_points['stage2_grasp_angle_class_acc/0_degree'] = (grasp_angle_class_pred==target_angles_cls)[objectness_mask.bool()].float().mean() 114 | acc_mask_15 = ((torch.abs(grasp_angle_class_pred-target_angles_cls)<=1) | (torch.abs(grasp_angle_class_pred-target_angles_cls)>=A-1)) 115 | end_points['stage2_grasp_angle_class_acc/15_degree'] = acc_mask_15[objectness_mask.bool()].float().mean() 116 | acc_mask_30 = ((torch.abs(grasp_angle_class_pred-target_angles_cls)<=2) | (torch.abs(grasp_angle_class_pred-target_angles_cls)>=A-2)) 117 | end_points['stage2_grasp_angle_class_acc/30_degree'] = acc_mask_30[objectness_mask.bool()].float().mean() 118 | 119 | 120 | # 3. depth cls loss 121 | target_depths_cls = torch.argmax(torch.max(batch_grasp_label, -2).values, -1) # B, N 122 | criterion_grasp_depth_class = nn.CrossEntropyLoss(reduction='none') 123 | grasp_depth_cls_score = end_points['grasp_depth_cls_pred'] # B, 4, N 124 | grasp_depth_cls_loss = criterion_grasp_depth_class(grasp_depth_cls_score, target_depths_cls) 125 | 126 | grasp_depth_cls_loss = grasp_depth_cls_loss[objectness_mask].mean() 127 | end_points['loss/stage2_grasp_depth_cls_loss'] = grasp_depth_cls_loss 128 | grasp_angle_class_pred = torch.argmax(grasp_depth_cls_score, 1) 129 | end_points['stage2_grasp_depth_cls_acc'] = (grasp_angle_class_pred==target_angles_cls)[objectness_mask.bool()].float().mean() 130 | 131 | 132 | # 4. width reg loss 133 | grasp_width_pred = end_points['grasp_width_pred'].transpose(1, 2).contiguous() # B, N, D 134 | grasp_width_loss = huber_loss((grasp_width_pred-target_widths)/GRASP_MAX_WIDTH, delta=1) 135 | objectness_mask4 = objectness_mask.unsqueeze(-1).expand_as(grasp_width_pred) 136 | grasp_width_loss = grasp_width_loss[objectness_mask4].mean() 137 | end_points['loss/stage2_grasp_width_loss'] = grasp_width_loss 138 | 139 | grasp_loss = grasp_score_loss + grasp_angle_class_loss + grasp_depth_cls_loss + 0.2*grasp_width_loss 140 | return grasp_loss, end_points 141 | 142 | 143 | def index_select(index, input): 144 | N = input.shape[0] 145 | output = [] 146 | for i in range(N): 147 | x = index[i] 148 | out = input[i][x] 149 | output.append(out) 150 | output = torch.stack(output) 151 | return output 152 | 153 | 154 | def pred_decode(end_points): 155 | batch_size = len(end_points['point_clouds']) 156 | grasp_preds = [] 157 | score_label = [] 158 | for i in range(batch_size): 159 | ## load predictions 160 | objectness_score = end_points['objectness_score'][i].float() 161 | grasp_score = end_points['grasp_score_pred'][i].float() # 48, N 162 | grasp_score = grasp_score.transpose(0, 1).view(-1, 12, 4) # N, 12, 4 163 | grasp_center = end_points['fp2_xyz'][i].float() 164 | approaching = -end_points['grasp_top_view_xyz'][i].float() 165 | grasp_angle_class_score = end_points['grasp_angle_cls_pred'][i] # 12, N 166 | grasp_depth_class_score = end_points['grasp_depth_cls_pred'][i] # 4, N 167 | grasp_width = 1.2 * end_points['grasp_width_pred'][i] # 4, N 168 | grasp_width = torch.clamp(grasp_width, min=0, max=GRASP_MAX_WIDTH) 169 | 170 | 171 | # grasp angle 172 | grasp_angle_class = torch.argmax(grasp_angle_class_score, 0) # N 173 | grasp_angle = grasp_angle_class.float() / 12 * np.pi 174 | # grasp score & width & tolerance 175 | #grasp_angle_class_ = grasp_angle_class.unsqueeze(0) # 1, N 176 | 177 | 178 | grasp_score = index_select(grasp_angle_class, grasp_score) # N, 4 179 | #grasp_width = torch.gather(grasp_width, 0, grasp_angle_class_).squeeze(0) 180 | 181 | 182 | # grasp depth 183 | grasp_depth_class = torch.argmax(grasp_depth_class_score, 0) # N 184 | grasp_depth = (grasp_depth_class.float() + 1) * 0.01 185 | # grasp score & angle & width & tolerance 186 | grasp_score = index_select(grasp_depth_class, grasp_score) # N 187 | # grasp_angle = torch.gather(grasp_angle, 1, grasp_depth_class) 188 | grasp_width = index_select(grasp_depth_class, grasp_width.transpose(0, 1)) # N 189 | 190 | 191 | ## slice preds by objectness 192 | objectness_pred = torch.argmax(objectness_score, 0) 193 | objectness_mask = (objectness_pred == 1) 194 | grasp_score = grasp_score[objectness_mask].unsqueeze(-1) 195 | grasp_width = grasp_width[objectness_mask].unsqueeze(-1) 196 | grasp_depth = grasp_depth[objectness_mask].unsqueeze(-1) 197 | approaching = approaching[objectness_mask] 198 | grasp_angle = grasp_angle[objectness_mask] 199 | 200 | # objectness_label = end_points['objectness_label'] 201 | # fp2_inds = end_points['fp2_inds'].long() 202 | # objectness_label = torch.gather(objectness_label, 1, fp2_inds).squeeze(0) 203 | # objectness_mask1 = (objectness_label==1) 204 | 205 | grasp_center = grasp_center[objectness_mask] 206 | 207 | # grasp_score = grasp_score * grasp_tolerance / GRASP_MAX_TOLERANCE 208 | 209 | ## convert to rotation matrix 210 | Ns = grasp_angle.size(0) 211 | approaching_ = approaching.view(Ns, 3) 212 | grasp_angle_ = grasp_angle.view(Ns) 213 | rotation_matrix = batch_viewpoint_params_to_matrix(approaching_, grasp_angle_) 214 | rotation_matrix = rotation_matrix.view(Ns, 9) 215 | 216 | # merge preds 217 | grasp_height = 0.02 * torch.ones_like(grasp_score) 218 | obj_ids = -1 * torch.ones_like(grasp_score) 219 | grasp_preds.append( 220 | torch.cat([grasp_score, grasp_width, grasp_height, grasp_depth, rotation_matrix, grasp_center, obj_ids], 221 | axis=-1)) 222 | 223 | grasp_label = end_points['batch_grasp_label'][i].float() # N, A, D 224 | grasp_label = index_select(grasp_angle_class, grasp_label) # N, 4 225 | grasp_label = index_select(grasp_depth_class, grasp_label) # N 226 | grasp_label = grasp_label[objectness_mask] 227 | score_label.append(grasp_label) 228 | return grasp_preds , score_label -------------------------------------------------------------------------------- /FGC_generate/score_gen_mp.py: -------------------------------------------------------------------------------- 1 | from vis_grasp import visObjGrasp, to_open3d_geometry_list 2 | import open3d as o3d 3 | import numpy as np 4 | import os 5 | from graspnetAPI.utils.utils import generate_views 6 | from graspnetAPI.utils.rotation import batch_viewpoint_params_to_matrix 7 | from tqdm import tqdm 8 | from numba import jit 9 | import multiprocessing as mp 10 | import time 11 | 12 | 13 | @jit(nopython=True) 14 | def cal_dist_nb(point1, point2, point3): 15 | ''' 16 | 17 | :param point1: (x1, y1, z1), the point 1 of line 18 | :param point2: (x2, y2, z2), the point 2 of line 19 | :param point3: (x3, y3, z3) 20 | v12 = point1-point2 21 | v13 = point1-point3 22 | distance = |v12×v13| / |v12| 23 | :return: dis 24 | 25 | ''' 26 | 27 | vec1 = point1 - point2 28 | vec2 = point1 - point3 29 | dis = abs(np.linalg.norm(np.cross(vec1, vec2))) / abs(np.linalg.norm(vec1)) 30 | dis13_left = np.linalg.norm(point1 - point3) 31 | dis23_right = np.linalg.norm(point2 - point3) 32 | if dis13_left <= dis23_right: 33 | # 0 means point3 close to left contact, 1 means point3 close to right contact 34 | dis = [dis, 0] 35 | else: 36 | dis = [dis, 1] 37 | return dis 38 | 39 | 40 | def load_grasp_label(root, obj_id): 41 | ''' 42 | score3: the score to estimate the flatness 43 | ''' 44 | obj_path = os.path.join(root, 'grasp_label', '{}_labels.npz'.format(str(obj_id).zfill(3))) 45 | label = np.load(obj_path) 46 | points = label['points'] 47 | offsets = label['offsets'] 48 | scores = label['scores'] 49 | collision = label['collision'] 50 | obj_idx = 0 51 | 52 | normal_path = os.path.join(root, 'normals_score', '{}_labels.npz'.format(str(obj_id).zfill(3))) 53 | normal_score = np.load(normal_path) 54 | normals = normal_score['normals'] 55 | score3 = normal_score['score3'] 56 | return points, offsets, scores, collision, obj_idx, normals, score3 57 | 58 | 59 | def get_grasp(root, obj_id): 60 | sampled_points, offsets, scores, _, obj_idx, normals, score3 = load_grasp_label(root, obj_id) 61 | num_views, num_angles, num_depths = 300, 12, 4 62 | 63 | point_inds = np.arange(sampled_points.shape[0]) 64 | template_views = generate_views(num_views) 65 | template_views = template_views[np.newaxis, :, np.newaxis, np.newaxis, :] 66 | template_views = np.tile(template_views, [1, 1, num_angles, num_depths, 1]) 67 | 68 | num_points = len(point_inds) 69 | target_points = sampled_points[:, np.newaxis, np.newaxis, np.newaxis, :] 70 | target_points = np.tile(target_points, [1, num_views, num_angles, num_depths, 1]) 71 | views = np.tile(template_views, [num_points, 1, 1, 1, 1]) 72 | angles = offsets[:, :, :, :, 0] 73 | depths = offsets[:, :, :, :, 1] 74 | widths = offsets[:, :, :, :, 2] 75 | 76 | mask1 = (scores > 0) 77 | mask1_idx = np.where(scores > 0) 78 | target_points = target_points[mask1] 79 | views = views[mask1] 80 | angles = angles[mask1] 81 | depths = depths[mask1] 82 | widths = widths[mask1] 83 | fric_coefs = scores[mask1] 84 | 85 | Rs = batch_viewpoint_params_to_matrix(-views, angles) 86 | 87 | num_grasp = widths.shape[0] 88 | scores = (1.1 - fric_coefs).reshape(-1, 1) 89 | widths = widths.reshape(-1, 1) 90 | heights = 0.02 * np.ones((num_grasp, 1)) 91 | depths = depths.reshape(-1, 1) 92 | rotations = Rs.reshape((-1, 9)) 93 | object_ids = obj_idx * np.ones((num_grasp, 1), dtype=np.int32) 94 | 95 | obj_grasp_array = np.hstack([scores, widths, heights, depths, rotations, target_points, object_ids]).astype( 96 | np.float32) 97 | 98 | return sampled_points, obj_grasp_array, mask1_idx, normals, score3 99 | 100 | 101 | @jit(nopython=True) 102 | def cal_dist(point1, point2, point3): 103 | ''' 104 | 105 | :param point1: (x1, y1, z1), the point 1 of line 106 | :param point2: (x2, y2, z2), the point 2 of line 107 | :param point3: (x3, y3, z3) 108 | v12 = point1-point2 109 | v13 = point1-point3 110 | distance = |v12×v13| / |v12| 111 | :return: dis 112 | 113 | ''' 114 | 115 | vec1 = point1-point2 116 | vec2 = point1-point3 117 | dis = abs(np.linalg.norm(np.cross(vec1, vec2)))/abs(np.linalg.norm(vec1)) 118 | dis13_left = np.linalg.norm(point1-point3) 119 | dis23_right = np.linalg.norm(point2-point3) 120 | if dis13_left <= dis23_right: 121 | # 0 means point3 close to left contact, 1 means point3 close to right contact 122 | dis = [dis, 0] 123 | else: 124 | dis = [dis, 1] 125 | return dis 126 | 127 | 128 | def collision(p1, p2, p3, p4): 129 | 130 | dis1 = np.linalg.norm(p1-p3) 131 | dis2 = np.linalg.norm(p2-p4) 132 | dec_ids = min(dis1, dis2) 133 | return dec_ids 134 | 135 | 136 | @jit(nopython=True) 137 | def cal_normal2lr(left_idx, right_idx, points, normals): 138 | ''' 139 | 140 | x = left_point_normal: on the object surface 141 | y = right_point_normal: on the object surface 142 | z = left_point-right_point: two contact point vector 143 | 144 | cos1 = |x.z / |x|.|z|| 145 | cos2 = |y.z / |y|.|z|| 146 | d = (cos1+cos2)/2 147 | 148 | :return: d 149 | ''' 150 | normals = np.asarray(normals, dtype=np.float32) 151 | left_point = points[left_idx] 152 | right_point = points[right_idx] 153 | 154 | left_normal = normals[left_idx] 155 | right_normal = normals[right_idx] 156 | 157 | vector_lr = left_point-right_point 158 | vector_lr_norm = np.linalg.norm(vector_lr) 159 | 160 | cos1 = abs(np.dot(left_normal, vector_lr)/(np.linalg.norm(left_normal)*vector_lr_norm)) 161 | cos2 = abs(np.dot(right_normal, vector_lr)/(np.linalg.norm(left_normal)*vector_lr_norm)) 162 | score4 = (cos1+cos2)/2 163 | return score4 164 | 165 | 166 | @jit(nopython=True) 167 | def init_contact(grasp_x): 168 | width = grasp_x[1] 169 | depth = grasp_x[3] 170 | rot = grasp_x[4:13].reshape((3, 3)) 171 | center = grasp_x[-4:-1] 172 | height = 0.004 173 | left_point = np.array([depth - height / 2, -width / 2, 0], dtype='float32') # 定义抓取接触点的初始点 174 | right_point = np.array([depth - height / 2, width / 2, 0], dtype='float32') 175 | 176 | left_contact = np.dot(rot, left_point.T).T + center # 得到旋转平移后的接触点 177 | right_contact = np.dot(rot, right_point.T).T + center 178 | return left_contact, right_contact 179 | 180 | 181 | def contact_find_solo(root): 182 | obj_pc, grasp, mask1_idx, normals, score3 = get_grasp(root) 183 | 184 | grasp = grasp[1] 185 | 186 | width = grasp[1] 187 | depth = grasp[3] 188 | rot = grasp[4:13].reshape((3, 3)) 189 | center = grasp[-4:-1] 190 | left_point, right_point = init_contact(width, depth) 191 | 192 | left_contact = np.dot(rot, left_point.T).T + center # 得到旋转平移后的接触点 193 | right_contact = np.dot(rot, right_point.T).T + center 194 | gravity_center = np.asarray([0, 0, 0]) 195 | 196 | pc_num = obj_pc.shape[0] 197 | dis = np.zeros((pc_num, 2)) 198 | for i in range(pc_num): 199 | point3 = obj_pc[i] 200 | dis_i = cal_dist_nb(left_contact, right_contact, point3) 201 | dis[i, :] = np.asarray(dis_i) 202 | 203 | min2max = np.argsort(dis[:, 0]) 204 | for i in min2max: 205 | if dis[i, 1] == 0: 206 | left_idx = i 207 | break 208 | 209 | for j in min2max: 210 | if dis[j, 1] == 1: 211 | right_idx = j 212 | break 213 | 214 | point_target_left = obj_pc[left_idx] 215 | point_target_right = obj_pc[right_idx] 216 | gravity_center_score = cal_dist(point_target_left, point_target_right, gravity_center) 217 | 218 | four_point = np.stack((point_target_left, point_target_right, left_contact, right_contact), axis=0) 219 | 220 | return obj_pc, grasp, four_point, gravity_center_score 221 | 222 | 223 | @jit(nopython=True) 224 | def contact_decision(obj_pc, left_contact, right_contact): 225 | pc_num = obj_pc.shape[0] 226 | dis = np.zeros((pc_num, 2)) 227 | for i in range(pc_num): 228 | point3 = obj_pc[i] 229 | dis_i = cal_dist(left_contact, right_contact, point3) 230 | dis[i, :] = np.asarray(dis_i) 231 | 232 | min2max = np.argsort(dis[:, 0]) 233 | 234 | for i in min2max: 235 | if dis[i, 1] == 0: 236 | left_idx = i 237 | break 238 | 239 | for j in min2max: 240 | if dis[j, 1] == 1: 241 | right_idx = j 242 | break 243 | 244 | return left_idx, right_idx 245 | 246 | 247 | def contact_find(root, obj_id): 248 | obj_pc, grasp, mask1_idx, normals, score3 = get_grasp(root, obj_id) 249 | 250 | grasp_num = grasp.shape[0] 251 | #grasp = grasp[1] 252 | 253 | four_point_all = [] 254 | gravity_score_all = [] 255 | flatness_score_all = [] 256 | consistency_score_all = [] 257 | collision_score_all = [] 258 | start = time.time() 259 | for x in tqdm(range(grasp_num), desc = 'Loading grasp...'): 260 | #start = time.time() 261 | 262 | grasp_x = grasp[x] 263 | left_contact, right_contact = init_contact(grasp_x) 264 | 265 | gravity_center = np.asarray([0, 0, 0]) 266 | 267 | left_idx, right_idx = contact_decision(obj_pc, left_contact, right_contact) 268 | #mid = time.time() 269 | 270 | point_target_left = obj_pc[left_idx] 271 | point_target_right = obj_pc[right_idx] 272 | 273 | Collision_perturbation_score = collision(point_target_left, point_target_right, left_contact, right_contact) 274 | 275 | gravity_center_score = cal_dist(point_target_left, point_target_right, gravity_center)[0] 276 | 277 | #four_point = np.stack((point_target_left, point_target_right, left_contact, right_contact), axis=0) 278 | 279 | flatness_score = (score3[left_idx]+score3[right_idx])/2 280 | 281 | score4 = cal_normal2lr(left_idx, right_idx, obj_pc, normals) 282 | 283 | #four_point_all.append(four_point) 284 | collision_score_all.append(Collision_perturbation_score) 285 | gravity_score_all.append(gravity_center_score) 286 | flatness_score_all.append(flatness_score) 287 | consistency_score_all.append(score4) 288 | 289 | #end = time.time() 290 | #print('for time', mid-start) 291 | #print('one frame time', end-start) 292 | #four_point_all = np.asarray(four_point_all) 293 | end = time.time() 294 | print('obj_id:', obj_id, 'finish time:', end-start) 295 | collision_score_all = np.asarray(collision_score_all) 296 | gravity_score_all = np.asarray(gravity_score_all) 297 | flatness_score_all = np.asarray(flatness_score_all) 298 | consistency_score_all = np.asarray(consistency_score_all) 299 | return obj_pc, grasp, mask1_idx, gravity_score_all, flatness_score_all, consistency_score_all, collision_score_all 300 | 301 | 302 | def score5_gen(root, obj_id): 303 | obj_pc, grasp, mask1_idx, normals, score3 = get_grasp(root, obj_id) 304 | 305 | grasp_num = grasp.shape[0] 306 | #grasp = grasp[1] 307 | 308 | four_point_all = [] 309 | gravity_score_all = [] 310 | flatness_score_all = [] 311 | consistency_score_all = [] 312 | collision_score_all = [] 313 | start = time.time() 314 | for x in tqdm(range(grasp_num), desc = 'Loading grasp...'): 315 | #start = time.time() 316 | 317 | grasp_x = grasp[x] 318 | left_contact, right_contact = init_contact(grasp_x) 319 | 320 | gravity_center = np.asarray([0, 0, 0]) 321 | 322 | left_idx, right_idx = contact_decision(obj_pc, left_contact, right_contact) 323 | #mid = time.time() 324 | 325 | point_target_left = obj_pc[left_idx] 326 | point_target_right = obj_pc[right_idx] 327 | 328 | Collision_perturbation_score = collision(point_target_left, point_target_right, left_contact, right_contact) 329 | 330 | collision_score_all.append(Collision_perturbation_score) 331 | end = time.time() 332 | print('obj_id:', obj_id, 'finish time:', end - start) 333 | collision_score_all = np.asarray(collision_score_all) 334 | return collision_score_all 335 | 336 | 337 | def vis_contact(self): 338 | obj_pc, grasp, four_point, score_gc = self.contact_find_solo() 339 | 340 | objp3d = o3d.geometry.PointCloud() 341 | objp3d.points = o3d.utility.Vector3dVector(obj_pc) 342 | objp3d.paint_uniform_color([0.3, 0.5, 0]) 343 | 344 | pc_target = o3d.geometry.PointCloud() 345 | pc_target.points = o3d.utility.Vector3dVector(four_point) 346 | pc_target.paint_uniform_color([0, 0, 1]) 347 | 348 | gg = to_open3d_geometry_list(grasp) 349 | 350 | o3d.visualization.draw_geometries([*gg, pc_target, objp3d], width=800, height=600, left=50, top=50) 351 | 352 | 353 | def save_new_score(root, obj_id): 354 | _, _, mask_idx, score2, score3, score4, score5 = contact_find(root, obj_id) 355 | savepath = os.path.join(root, 'new_score', '{}_labels.npz'.format(str(obj_id).zfill(3))) 356 | np.savez(savepath, mask_idx=mask_idx, score2=score2, score3=score3, score4=score4, score5=score5) 357 | 358 | 359 | def save_score5(root, obj_id): 360 | score5 = score5_gen(root, obj_id) 361 | score5_norm = (score5 - np.min(score5)) / (np.max(score5) - np.min(score5)) 362 | savepath = os.path.join(root, 'score5', '{}_labels.npz'.format(str(obj_id).zfill(3))) 363 | np.savez(savepath, score5=score5_norm) 364 | 365 | if __name__=='__main__': 366 | root = '../grasp_data' 367 | pool_size = 32 368 | for obj_id in range(88): 369 | p = mp.Process(target=save_score5, args=(root, obj_id)) 370 | p.start() 371 | 372 | --------------------------------------------------------------------------------