├── .flake8 ├── .gitignore ├── LICENSE ├── README.md ├── pyproject.toml ├── requirements.txt ├── setup.py └── vision3d ├── __init__.py ├── array_ops ├── __init__.py ├── ball_query.py ├── deformation.py ├── depth_image.py ├── furthest_point_sample.py ├── graph_pyramid.py ├── grid_reduce.py ├── grid_subsample.py ├── keypoint_utils.py ├── knn.py ├── knn_interpolate.py ├── metrics.py ├── mutual_select.py ├── point_cloud_utils.py ├── procrustes.py ├── radius_nms.py ├── radius_search.py ├── ray_utils.py ├── registration_utils.py ├── se3.py └── so3.py ├── csrc ├── .clang-format ├── common │ ├── cuda_helper.h │ ├── numpy_helper.h │ └── torch_helper.h ├── cpu │ ├── cloud │ │ ├── cloud.cpp │ │ └── cloud.h │ ├── deformation_graph │ │ ├── deformation_graph.cpp │ │ └── deformation_graph.h │ ├── graph_proc │ │ ├── graph_proc.cpp │ │ └── graph_proc.h │ ├── grid_subsampling │ │ ├── grid_subsampling.cpp │ │ └── grid_subsampling.h │ ├── node_sampling │ │ ├── node_sampling.cpp │ │ └── node_sampling.h │ ├── radius_neighbors │ │ ├── radius_neighbors.cpp │ │ └── radius_neighbors.h │ └── voxel │ │ ├── voxel.cpp │ │ └── voxel.h ├── cuda │ ├── ball_query │ │ ├── ball_query.cpp │ │ ├── ball_query.h │ │ ├── ball_query_cuda.cu │ │ └── ball_query_cuda.cuh │ ├── furthest_point_sample │ │ ├── furthest_point_sample.cpp │ │ ├── furthest_point_sample.h │ │ ├── furthest_point_sample_cuda.cu │ │ └── furthest_point_sample_cuda.cuh │ ├── gather │ │ ├── gather.cpp │ │ ├── gather.h │ │ ├── gather_cuda.cu │ │ └── gather_cuda.cuh │ ├── group_gather │ │ ├── group_gather.cpp │ │ ├── group_gather.h │ │ ├── group_gather_cuda.cu │ │ └── group_gather_cuda.cuh │ ├── knn_points │ │ ├── knn_points.cpp │ │ ├── knn_points.h │ │ ├── knn_points_cuda.cu │ │ └── knn_points_cuda.cuh │ ├── three_interpolate │ │ ├── three_interpolate.cpp │ │ ├── three_interpolate.h │ │ ├── three_interpolate_cuda.cu │ │ └── three_interpolate_cuda.cuh │ └── three_nn │ │ ├── three_nn.cpp │ │ ├── three_nn.h │ │ ├── three_nn_cuda.cu │ │ └── three_nn_cuda.cuh ├── external │ ├── eigen3 │ │ └── Eigen │ │ │ ├── Cholesky │ │ │ ├── CholmodSupport │ │ │ ├── Core │ │ │ ├── Dense │ │ │ ├── Eigen │ │ │ ├── Eigenvalues │ │ │ ├── Geometry │ │ │ ├── Householder │ │ │ ├── IterativeLinearSolvers │ │ │ ├── Jacobi │ │ │ ├── LU │ │ │ ├── MetisSupport │ │ │ ├── OrderingMethods │ │ │ ├── PaStiXSupport │ │ │ ├── PardisoSupport │ │ │ ├── QR │ │ │ ├── QtAlignedMalloc │ │ │ ├── SPQRSupport │ │ │ ├── SVD │ │ │ ├── Sparse │ │ │ ├── SparseCholesky │ │ │ ├── SparseCore │ │ │ ├── SparseLU │ │ │ ├── SparseQR │ │ │ ├── StdDeque │ │ │ ├── StdList │ │ │ ├── StdVector │ │ │ ├── SuperLUSupport │ │ │ ├── UmfPackSupport │ │ │ └── src │ │ │ ├── Cholesky │ │ │ ├── LDLT.h │ │ │ ├── LLT.h │ │ │ └── LLT_LAPACKE.h │ │ │ ├── CholmodSupport │ │ │ └── CholmodSupport.h │ │ │ ├── Core │ │ │ ├── Array.h │ │ │ ├── ArrayBase.h │ │ │ ├── ArrayWrapper.h │ │ │ ├── Assign.h │ │ │ ├── AssignEvaluator.h │ │ │ ├── Assign_MKL.h │ │ │ ├── BandMatrix.h │ │ │ ├── Block.h │ │ │ ├── BooleanRedux.h │ │ │ ├── CommaInitializer.h │ │ │ ├── ConditionEstimator.h │ │ │ ├── CoreEvaluators.h │ │ │ ├── CoreIterators.h │ │ │ ├── CwiseBinaryOp.h │ │ │ ├── CwiseNullaryOp.h │ │ │ ├── CwiseTernaryOp.h │ │ │ ├── CwiseUnaryOp.h │ │ │ ├── CwiseUnaryView.h │ │ │ ├── DenseBase.h │ │ │ ├── DenseCoeffsBase.h │ │ │ ├── DenseStorage.h │ │ │ ├── Diagonal.h │ │ │ ├── DiagonalMatrix.h │ │ │ ├── DiagonalProduct.h │ │ │ ├── Dot.h │ │ │ ├── EigenBase.h │ │ │ ├── ForceAlignedAccess.h │ │ │ ├── Fuzzy.h │ │ │ ├── GeneralProduct.h │ │ │ ├── GenericPacketMath.h │ │ │ ├── GlobalFunctions.h │ │ │ ├── IO.h │ │ │ ├── Inverse.h │ │ │ ├── Map.h │ │ │ ├── MapBase.h │ │ │ ├── MathFunctions.h │ │ │ ├── MathFunctionsImpl.h │ │ │ ├── Matrix.h │ │ │ ├── MatrixBase.h │ │ │ ├── NestByValue.h │ │ │ ├── NoAlias.h │ │ │ ├── NumTraits.h │ │ │ ├── PermutationMatrix.h │ │ │ ├── PlainObjectBase.h │ │ │ ├── Product.h │ │ │ ├── ProductEvaluators.h │ │ │ ├── Random.h │ │ │ ├── Redux.h │ │ │ ├── Ref.h │ │ │ ├── Replicate.h │ │ │ ├── ReturnByValue.h │ │ │ ├── Reverse.h │ │ │ ├── Select.h │ │ │ ├── SelfAdjointView.h │ │ │ ├── SelfCwiseBinaryOp.h │ │ │ ├── Solve.h │ │ │ ├── SolveTriangular.h │ │ │ ├── SolverBase.h │ │ │ ├── StableNorm.h │ │ │ ├── Stride.h │ │ │ ├── Swap.h │ │ │ ├── Transpose.h │ │ │ ├── Transpositions.h │ │ │ ├── TriangularMatrix.h │ │ │ ├── VectorBlock.h │ │ │ ├── VectorwiseOp.h │ │ │ ├── Visitor.h │ │ │ ├── arch │ │ │ │ ├── AVX │ │ │ │ │ ├── Complex.h │ │ │ │ │ ├── MathFunctions.h │ │ │ │ │ ├── PacketMath.h │ │ │ │ │ └── TypeCasting.h │ │ │ │ ├── AVX512 │ │ │ │ │ ├── MathFunctions.h │ │ │ │ │ └── PacketMath.h │ │ │ │ ├── AltiVec │ │ │ │ │ ├── Complex.h │ │ │ │ │ ├── MathFunctions.h │ │ │ │ │ └── PacketMath.h │ │ │ │ ├── CUDA │ │ │ │ │ ├── Complex.h │ │ │ │ │ ├── Half.h │ │ │ │ │ ├── MathFunctions.h │ │ │ │ │ ├── PacketMath.h │ │ │ │ │ ├── PacketMathHalf.h │ │ │ │ │ └── TypeCasting.h │ │ │ │ ├── Default │ │ │ │ │ ├── ConjHelper.h │ │ │ │ │ └── Settings.h │ │ │ │ ├── NEON │ │ │ │ │ ├── Complex.h │ │ │ │ │ ├── MathFunctions.h │ │ │ │ │ └── PacketMath.h │ │ │ │ ├── SSE │ │ │ │ │ ├── Complex.h │ │ │ │ │ ├── MathFunctions.h │ │ │ │ │ ├── PacketMath.h │ │ │ │ │ └── TypeCasting.h │ │ │ │ └── ZVector │ │ │ │ │ ├── Complex.h │ │ │ │ │ ├── MathFunctions.h │ │ │ │ │ └── PacketMath.h │ │ │ ├── functors │ │ │ │ ├── AssignmentFunctors.h │ │ │ │ ├── BinaryFunctors.h │ │ │ │ ├── NullaryFunctors.h │ │ │ │ ├── StlFunctors.h │ │ │ │ ├── TernaryFunctors.h │ │ │ │ └── UnaryFunctors.h │ │ │ ├── products │ │ │ │ ├── GeneralBlockPanelKernel.h │ │ │ │ ├── GeneralMatrixMatrix.h │ │ │ │ ├── GeneralMatrixMatrixTriangular.h │ │ │ │ ├── GeneralMatrixMatrixTriangular_BLAS.h │ │ │ │ ├── GeneralMatrixMatrix_BLAS.h │ │ │ │ ├── GeneralMatrixVector.h │ │ │ │ ├── GeneralMatrixVector_BLAS.h │ │ │ │ ├── Parallelizer.h │ │ │ │ ├── SelfadjointMatrixMatrix.h │ │ │ │ ├── SelfadjointMatrixMatrix_BLAS.h │ │ │ │ ├── SelfadjointMatrixVector.h │ │ │ │ ├── SelfadjointMatrixVector_BLAS.h │ │ │ │ ├── SelfadjointProduct.h │ │ │ │ ├── SelfadjointRank2Update.h │ │ │ │ ├── TriangularMatrixMatrix.h │ │ │ │ ├── TriangularMatrixMatrix_BLAS.h │ │ │ │ ├── TriangularMatrixVector.h │ │ │ │ ├── TriangularMatrixVector_BLAS.h │ │ │ │ ├── TriangularSolverMatrix.h │ │ │ │ ├── TriangularSolverMatrix_BLAS.h │ │ │ │ └── TriangularSolverVector.h │ │ │ └── util │ │ │ │ ├── BlasUtil.h │ │ │ │ ├── Constants.h │ │ │ │ ├── DisableStupidWarnings.h │ │ │ │ ├── ForwardDeclarations.h │ │ │ │ ├── MKL_support.h │ │ │ │ ├── Macros.h │ │ │ │ ├── Memory.h │ │ │ │ ├── Meta.h │ │ │ │ ├── NonMPL2.h │ │ │ │ ├── ReenableStupidWarnings.h │ │ │ │ ├── StaticAssert.h │ │ │ │ └── XprHelper.h │ │ │ ├── Eigenvalues │ │ │ ├── ComplexEigenSolver.h │ │ │ ├── ComplexSchur.h │ │ │ ├── ComplexSchur_LAPACKE.h │ │ │ ├── EigenSolver.h │ │ │ ├── GeneralizedEigenSolver.h │ │ │ ├── GeneralizedSelfAdjointEigenSolver.h │ │ │ ├── HessenbergDecomposition.h │ │ │ ├── MatrixBaseEigenvalues.h │ │ │ ├── RealQZ.h │ │ │ ├── RealSchur.h │ │ │ ├── RealSchur_LAPACKE.h │ │ │ ├── SelfAdjointEigenSolver.h │ │ │ ├── SelfAdjointEigenSolver_LAPACKE.h │ │ │ └── Tridiagonalization.h │ │ │ ├── Geometry │ │ │ ├── AlignedBox.h │ │ │ ├── AngleAxis.h │ │ │ ├── EulerAngles.h │ │ │ ├── Homogeneous.h │ │ │ ├── Hyperplane.h │ │ │ ├── OrthoMethods.h │ │ │ ├── ParametrizedLine.h │ │ │ ├── Quaternion.h │ │ │ ├── Rotation2D.h │ │ │ ├── RotationBase.h │ │ │ ├── Scaling.h │ │ │ ├── Transform.h │ │ │ ├── Translation.h │ │ │ ├── Umeyama.h │ │ │ └── arch │ │ │ │ └── Geometry_SSE.h │ │ │ ├── Householder │ │ │ ├── BlockHouseholder.h │ │ │ ├── Householder.h │ │ │ └── HouseholderSequence.h │ │ │ ├── IterativeLinearSolvers │ │ │ ├── BasicPreconditioners.h │ │ │ ├── BiCGSTAB.h │ │ │ ├── ConjugateGradient.h │ │ │ ├── IncompleteCholesky.h │ │ │ ├── IncompleteLUT.h │ │ │ ├── IterativeSolverBase.h │ │ │ ├── LeastSquareConjugateGradient.h │ │ │ └── SolveWithGuess.h │ │ │ ├── Jacobi │ │ │ └── Jacobi.h │ │ │ ├── LU │ │ │ ├── Determinant.h │ │ │ ├── FullPivLU.h │ │ │ ├── InverseImpl.h │ │ │ ├── PartialPivLU.h │ │ │ ├── PartialPivLU_LAPACKE.h │ │ │ └── arch │ │ │ │ └── Inverse_SSE.h │ │ │ ├── MetisSupport │ │ │ └── MetisSupport.h │ │ │ ├── OrderingMethods │ │ │ ├── Amd.h │ │ │ ├── Eigen_Colamd.h │ │ │ └── Ordering.h │ │ │ ├── PaStiXSupport │ │ │ └── PaStiXSupport.h │ │ │ ├── PardisoSupport │ │ │ └── PardisoSupport.h │ │ │ ├── QR │ │ │ ├── ColPivHouseholderQR.h │ │ │ ├── ColPivHouseholderQR_LAPACKE.h │ │ │ ├── CompleteOrthogonalDecomposition.h │ │ │ ├── FullPivHouseholderQR.h │ │ │ ├── HouseholderQR.h │ │ │ └── HouseholderQR_LAPACKE.h │ │ │ ├── SPQRSupport │ │ │ └── SuiteSparseQRSupport.h │ │ │ ├── SVD │ │ │ ├── BDCSVD.h │ │ │ ├── JacobiSVD.h │ │ │ ├── JacobiSVD_LAPACKE.h │ │ │ ├── SVDBase.h │ │ │ └── UpperBidiagonalization.h │ │ │ ├── SparseCholesky │ │ │ ├── SimplicialCholesky.h │ │ │ └── SimplicialCholesky_impl.h │ │ │ ├── SparseCore │ │ │ ├── AmbiVector.h │ │ │ ├── CompressedStorage.h │ │ │ ├── ConservativeSparseSparseProduct.h │ │ │ ├── MappedSparseMatrix.h │ │ │ ├── SparseAssign.h │ │ │ ├── SparseBlock.h │ │ │ ├── SparseColEtree.h │ │ │ ├── SparseCompressedBase.h │ │ │ ├── SparseCwiseBinaryOp.h │ │ │ ├── SparseCwiseUnaryOp.h │ │ │ ├── SparseDenseProduct.h │ │ │ ├── SparseDiagonalProduct.h │ │ │ ├── SparseDot.h │ │ │ ├── SparseFuzzy.h │ │ │ ├── SparseMap.h │ │ │ ├── SparseMatrix.h │ │ │ ├── SparseMatrixBase.h │ │ │ ├── SparsePermutation.h │ │ │ ├── SparseProduct.h │ │ │ ├── SparseRedux.h │ │ │ ├── SparseRef.h │ │ │ ├── SparseSelfAdjointView.h │ │ │ ├── SparseSolverBase.h │ │ │ ├── SparseSparseProductWithPruning.h │ │ │ ├── SparseTranspose.h │ │ │ ├── SparseTriangularView.h │ │ │ ├── SparseUtil.h │ │ │ ├── SparseVector.h │ │ │ ├── SparseView.h │ │ │ └── TriangularSolver.h │ │ │ ├── SparseLU │ │ │ ├── SparseLU.h │ │ │ ├── SparseLUImpl.h │ │ │ ├── SparseLU_Memory.h │ │ │ ├── SparseLU_Structs.h │ │ │ ├── SparseLU_SupernodalMatrix.h │ │ │ ├── SparseLU_Utils.h │ │ │ ├── SparseLU_column_bmod.h │ │ │ ├── SparseLU_column_dfs.h │ │ │ ├── SparseLU_copy_to_ucol.h │ │ │ ├── SparseLU_gemm_kernel.h │ │ │ ├── SparseLU_heap_relax_snode.h │ │ │ ├── SparseLU_kernel_bmod.h │ │ │ ├── SparseLU_panel_bmod.h │ │ │ ├── SparseLU_panel_dfs.h │ │ │ ├── SparseLU_pivotL.h │ │ │ ├── SparseLU_pruneL.h │ │ │ └── SparseLU_relax_snode.h │ │ │ ├── SparseQR │ │ │ └── SparseQR.h │ │ │ ├── StlSupport │ │ │ ├── StdDeque.h │ │ │ ├── StdList.h │ │ │ ├── StdVector.h │ │ │ └── details.h │ │ │ ├── SuperLUSupport │ │ │ └── SuperLUSupport.h │ │ │ ├── UmfPackSupport │ │ │ └── UmfPackSupport.h │ │ │ ├── misc │ │ │ ├── Image.h │ │ │ ├── Kernel.h │ │ │ ├── RealSvd2x2.h │ │ │ ├── blas.h │ │ │ ├── lapack.h │ │ │ ├── lapacke.h │ │ │ └── lapacke_mangling.h │ │ │ └── plugins │ │ │ ├── ArrayCwiseBinaryOps.h │ │ │ ├── ArrayCwiseUnaryOps.h │ │ │ ├── BlockMethods.h │ │ │ ├── CommonCwiseBinaryOps.h │ │ │ ├── CommonCwiseUnaryOps.h │ │ │ ├── MatrixCwiseBinaryOps.h │ │ │ └── MatrixCwiseUnaryOps.h │ └── nanoflann │ │ └── nanoflann.hpp └── pybind.cpp ├── datasets ├── README.md ├── __init__.py ├── classification │ ├── __init__.py │ └── modelnet40.py ├── completion │ ├── __init__.py │ └── modelnet.py ├── registration │ ├── __init__.py │ ├── cape │ │ ├── __init__.py │ │ └── cape.py │ ├── deepdeform │ │ ├── __init__.py │ │ └── deepdeform.py │ ├── fourdmatch │ │ ├── __init__.py │ │ ├── fourdmatch.py │ │ └── utils.py │ ├── kitti │ │ ├── __init__.py │ │ └── odometry_kitti.py │ ├── modelnet │ │ ├── __init__.py │ │ └── modelnet.py │ ├── redwood │ │ ├── __init__.py │ │ └── redwood.py │ ├── rgbdscenes │ │ ├── __init__.py │ │ └── rgbdscenes.py │ ├── scannet_urr │ │ ├── __init__.py │ │ └── scannet_urr.py │ ├── sevenscenes │ │ ├── __init__.py │ │ └── sevenscenes_hard.py │ ├── shapenet │ │ ├── __init__.py │ │ └── shapenet.py │ ├── threedmatch │ │ ├── __init__.py │ │ ├── threedmatch.py │ │ ├── threedmatch_rgb.py │ │ └── threedmatch_utils.py │ └── threedmatch_urr │ │ ├── __init__.py │ │ └── threedmatch_urr.py └── segmentation │ ├── __init__.py │ ├── s3dis.py │ ├── s3dis_voxelized.py │ ├── shapenetpart.py │ └── utils.py ├── engine ├── __init__.py ├── base_tester.py ├── base_trainer.py ├── batch_tester.py ├── checkpoint.py ├── context_manager.py ├── epoch_based_trainer.py ├── iter_based_trainer.py ├── single_tester.py └── utils.py ├── layers ├── __init__.py ├── basic_layers │ ├── __init__.py │ ├── builder.py │ ├── depthwise_conv.py │ ├── monte_carlo_dropout.py │ ├── norm.py │ ├── separable_conv.py │ └── utils.py ├── conv_block.py ├── edge_conv.py ├── embedding.py ├── feature_propagate.py ├── kpconv.py ├── kpconv_utils │ ├── __init__.py │ ├── dispositions │ │ └── k_015_center_3D.ply │ └── kernel_points.py ├── nonrigid_icp.py ├── pointnet.py ├── pointnet2.py ├── residual_block.py ├── sinkhorn.py ├── transformer.py ├── unary_block.py ├── vector_neurons.py ├── weighted_procrustes.py └── xconv.py ├── loss ├── __init__.py ├── as_rigid_as_possible_loss.py ├── chamfer_distance.py ├── circle_loss.py ├── focal_loss.py ├── hardest_contrastive_loss.py ├── orthogonal_loss.py ├── smooth_ce_loss.py ├── transformation_loss.py └── weighted_bce_loss.py ├── metrics ├── __init__.py ├── accuracy.py ├── iou.py └── part_iou.py ├── models ├── __init__.py ├── backbone │ ├── __init__.py │ └── kpconv_fpn.py ├── cofinet │ ├── __init__.py │ ├── node_matching.py │ └── node_proposal.py ├── contrib │ └── __init__.py ├── d3feat │ ├── __init__.py │ ├── detection_score.py │ └── loss.py ├── fcgf │ ├── __init__.py │ └── modules.py ├── foldingnet │ ├── __init__.py │ └── foldingnet.py ├── geotransformer │ ├── __init__.py │ ├── adaptive_superpoint_matching.py │ ├── geometric_structure_embedding.py │ ├── geometric_transformer.py │ ├── geometric_transformer_lite.py │ ├── local_global_registration.py │ ├── local_global_registration_duplicate_removal.py │ ├── local_global_registration_threshold.py │ ├── point_matching.py │ ├── spatial_consistency_filtering.py │ ├── superpoint_matching_global_topk.py │ ├── superpoint_matching_mutual_topk.py │ └── superpoint_proposal.py ├── graph_pyramid │ ├── __init__.py │ ├── calibrate_neighbors.py │ └── graph_pyramid.py ├── pat │ ├── __init__.py │ ├── functional.py │ └── modules.py ├── point_transformer │ ├── __init__.py │ └── modules.py ├── pointnet │ ├── __init__.py │ └── modules.py ├── pointnet2 │ ├── __init__.py │ └── modules.py ├── predator │ ├── __init__.py │ ├── loss.py │ └── modules.py ├── pri3d │ ├── __init__.py │ ├── resnet.py │ └── resunet.py ├── siamese_transformer │ ├── __init__.py │ └── vanilla_siamese_transformer.py └── urr │ ├── __init__.py │ ├── matching.py │ ├── registration.py │ └── render.py ├── ops ├── __init__.py ├── back_project.py ├── ball_query.py ├── contrib │ └── __init__.py ├── conversion.py ├── correspondences.py ├── cosine_similarity.py ├── deformation_graph.py ├── dual_softmax.py ├── eigenvector.py ├── embedded_deformation.py ├── furthest_point_sample.py ├── gather.py ├── grid_subsample.py ├── group_gather.py ├── index_select.py ├── knn.py ├── knn_interpolate.py ├── knn_points.py ├── local_reference_frame.py ├── masked_ops.py ├── matching.py ├── meshgrid.py ├── metrics.py ├── mutual_topk_select.py ├── nearest_interpolate.py ├── nonrigid_icp_adam.py ├── normal_estimation.py ├── numeric.py ├── pairwise_distance.py ├── point_cloud_partition.py ├── point_pair_feature.py ├── pooling.py ├── radius_search.py ├── random_point_sample.py ├── random_sample.py ├── render.py ├── sample_pdf.py ├── se3.py ├── so3.py ├── spatial_consistency.py ├── three_interpolate.py ├── three_nn.py ├── vector_angle.py ├── volume_render.py ├── voxelize.py └── weighted_procrustes.py └── utils ├── __init__.py ├── average_meter.py ├── collate.py ├── dataloader.py ├── distributed.py ├── io.py ├── logger.py ├── misc.py ├── open3d.py ├── opencv.py ├── optimizer.py ├── parser.py ├── profiling.py ├── summary_board.py ├── tensor.py ├── timer.py └── visualization.py /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | max-line-length = 120 3 | ignore = E501,W503 4 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # exclude package building files 2 | vision3d.egg-info 3 | vision3d/dist 4 | **/__pycache__ 5 | **/*.so 6 | **/build 7 | **/.idea 8 | **/.vscode 9 | 10 | # exclude project files 11 | trash 12 | output 13 | experiments 14 | pretrained 15 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Zheng Qin 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Vision3D: A 3D Vision Library built with PyTorch 2 | 3 | Vision3D is a 3D vision library built with PyTorch, in order to reproduce the state-of-the-art 3D vision models. 4 | 5 | ## Features 6 | 7 | 1. State-of-the-art models: 8 | 1. Point cloud processing: 9 | 1. PointNet series: PointNet, PointNet++. 10 | 2. DGCNN. 11 | 3. PointCNN. 12 | 4. Point Transformer. 13 | 5. KPConv. 14 | 2. Point cloud registration: 15 | 1. FCGF. 16 | 2. D3Feat. 17 | 3. Predator. 18 | 4. GeoTransformer 19 | 5. RPM-Net. 20 | 2. Multiple 3D vision tasks: 21 | 1. Object Classification. 22 | 2. Semantic Segmentation. 23 | 3. Registration. 24 | 4. Completion. 25 | 3. Multiple datasets: 26 | 1. Object Classification: ModelNet40. 27 | 2. Part Semantic Segmentation: ShapeNetPart. 28 | 3. Semantic Segmentation: S3DIS. 29 | 4. Point Cloud Registration: ModelNet40, 3DMatch, KITTI odometry, 4DMatch, DeepDeform, CAPE, Augmented ICL-NUIM. 30 | 4. Multi-GPU Training with DDP. 31 | 32 | ## Installation 33 | 34 | Vision3D is tested on Python 3.8, PyTorch 1.13.1, Ubuntu 22.04, GCC 11.3 and CUDA 11.7, but it should work with other configurations. Currently, Vision3d only support training and testing on GPUs. 35 | 36 | Install Vision3D with the following command: 37 | 38 | ```bash 39 | python setup.py develop 40 | ``` 41 | 42 | All requirements will be automatically installed. 43 | 44 | ## Acknowledgements 45 | 46 | - [PointNet](https://github.com/charlesq34/pointnet) 47 | - [PointNet2](https://github.com/charlesq34/pointnet2) 48 | - [DGCNN](https://github.com/WangYueFt/dgcnn) 49 | - [PointCNN](https://github.com/yangyanli/PointCNN) 50 | - [MMCV](https://github.com/open-mmlab/mmcv) 51 | - [KPConv](https://github.com/HuguesTHOMAS/KPConv-PyTorch) 52 | - [FCGF](https://github.com/chrischoy/FCGF) 53 | - [D3Feat](https://github.com/XuyangBai/D3Feat.pytorch) 54 | - [Predator](https://github.com/prs-eth/OverlapPredator) 55 | - [RPM-Net](https://github.com/yewzijian/RPMNet) 56 | - [UnsupervisedR&R](https://github.com/mbanani/unsupervisedRR) 57 | - [CoFiNet](https://github.com/haoyu94/Coarse-to-fine-correspondences) 58 | - [VectorNeuron](https://github.com/FlyingGiraffe/vnn) 59 | - And more. 60 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.black] 2 | line-length = 120 3 | 4 | [tool.isort] 5 | profile = "black" 6 | line_length = 120 7 | known_first_party = ["vision3d"] -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch>=1.11.0 2 | numpy 3 | scipy 4 | tqdm 5 | loguru 6 | easydict 7 | h5py 8 | pykeops 9 | torch_scatter 10 | open3d==0.11.2 11 | scikit-learn 12 | scikit-image 13 | einops 14 | tensorboard 15 | IPython 16 | ipdb 17 | opencv-python -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | from glob import glob 3 | 4 | from setuptools import setup 5 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 6 | 7 | 8 | extension_dir = "vision3d/csrc" 9 | 10 | 11 | def get_sources(): 12 | sources = [] 13 | sources += glob(osp.join(extension_dir, "cpu", "*", "*.cpp")) 14 | sources += glob(osp.join(extension_dir, "cuda", "*", "*.cpp")) 15 | sources += glob(osp.join(extension_dir, "cuda", "*", "*.cu")) 16 | sources += glob(osp.join(extension_dir, "pybind.cpp")) 17 | return sources 18 | 19 | 20 | def get_include_dirs(): 21 | include_dirs = [] 22 | include_dirs.append(osp.abspath(osp.join(extension_dir, "external", "eigen3"))) 23 | include_dirs.append(osp.abspath(osp.join(extension_dir, "external", "nanoflann"))) 24 | include_dirs.append(osp.abspath(osp.join(extension_dir, "common"))) 25 | return include_dirs 26 | 27 | 28 | def get_requirements(): 29 | with open("requirements.txt", "r") as f: 30 | lines = f.readlines() 31 | requirements = [line.strip() for line in lines] 32 | return requirements 33 | 34 | 35 | setup( 36 | name="vision3d", 37 | version="2.0.1", 38 | install_requires=get_requirements(), 39 | ext_modules=[ 40 | CUDAExtension( 41 | name="vision3d.ext", 42 | sources=get_sources(), 43 | include_dirs=get_include_dirs(), 44 | ), 45 | ], 46 | cmdclass={"build_ext": BuildExtension}, 47 | ) 48 | -------------------------------------------------------------------------------- /vision3d/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qinzheng93/vision3d/7c61728eeeac76c9febf4f70c399358ce1a81b25/vision3d/__init__.py -------------------------------------------------------------------------------- /vision3d/array_ops/ball_query.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | from numpy import ndarray 4 | from scipy.spatial.ckdtree import cKDTree 5 | 6 | 7 | def ball_query(q_points: ndarray, s_points: ndarray, radius: float) -> List[ndarray]: 8 | s_tree = cKDTree(s_points) 9 | indices_list = s_tree.query_ball_point(q_points, radius) 10 | return indices_list 11 | -------------------------------------------------------------------------------- /vision3d/array_ops/furthest_point_sample.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import numpy as np 4 | from numpy import ndarray 5 | 6 | from vision3d.utils.misc import load_ext 7 | 8 | ext_module = load_ext("vision3d.ext", ["sample_nodes_with_fps"]) 9 | 10 | 11 | def furthest_point_sample( 12 | points: ndarray, min_distance: Optional[float] = None, num_samples: Optional[int] = None 13 | ) -> ndarray: 14 | if min_distance is None: 15 | min_distance = -1.0 16 | if num_samples is None: 17 | num_samples = -1 18 | node_indices = ext_module.sample_nodes_with_fps(points, min_distance, num_samples) 19 | return node_indices 20 | 21 | 22 | def native_furthest_point_sample( 23 | points: ndarray, min_distance: Optional[float] = None, num_samples: Optional[int] = None 24 | ) -> ndarray: 25 | assert min_distance is not None or num_samples is not None, "Both 'min_distance' and 'num_samples' are None." 26 | 27 | num_points = points.shape[0] 28 | selected_indices = [] 29 | unselected_indices = list(range(num_points)) 30 | distances = np.full(shape=(num_points,), fill_value=np.inf) 31 | 32 | best_index = 0 33 | while True: 34 | sel_index = unselected_indices[best_index] 35 | selected_indices.append(sel_index) 36 | unselected_indices.pop(best_index) 37 | 38 | best_index = -1 39 | best_distance = 0.0 40 | indices_to_remove = [] 41 | for index, point_index in enumerate(unselected_indices): 42 | distance = np.linalg.norm(points[point_index] - points[sel_index]) 43 | distance = min(distances[point_index], distance) 44 | distances[point_index] = distance 45 | if distance > best_distance: 46 | best_index = index 47 | best_distance = distance 48 | if distance < min_distance: 49 | indices_to_remove.append(index) 50 | 51 | if indices_to_remove: 52 | indices_to_remove = reversed(indices_to_remove) 53 | for index in indices_to_remove: 54 | unselected_indices.pop(index) 55 | if index < best_index: 56 | best_index -= 1 57 | 58 | if min_distance is not None and best_distance < min_distance: 59 | break 60 | 61 | if num_samples is not None and len(selected_indices) >= num_samples: 62 | break 63 | 64 | sel_indices = np.asarray(selected_indices, dtype=np.int64) 65 | 66 | return sel_indices 67 | -------------------------------------------------------------------------------- /vision3d/array_ops/graph_pyramid.py: -------------------------------------------------------------------------------- 1 | from typing import List, Dict 2 | 3 | from numpy import ndarray 4 | 5 | from .grid_subsample import grid_subsample_pack_mode 6 | from .radius_search import radius_search_pack_mode 7 | 8 | 9 | def build_grid_and_radius_graph_pyramid_pack_mode( 10 | points: ndarray, 11 | lengths: ndarray, 12 | num_stages: int, 13 | voxel_size: float, 14 | radius: float, 15 | neighbor_limits: List[int], 16 | ) -> Dict: 17 | """Build graph pyramid with grid subsampling and radius searching in pack mode.""" 18 | assert num_stages is not None 19 | assert voxel_size is not None 20 | assert radius is not None 21 | assert neighbor_limits is not None 22 | assert num_stages == len(neighbor_limits) 23 | 24 | points_list = [] 25 | lengths_list = [] 26 | neighbors_list = [] 27 | subsampling_list = [] 28 | upsampling_list = [] 29 | 30 | # grid subsampling 31 | for i in range(num_stages): 32 | if i > 0: 33 | points, lengths = grid_subsample_pack_mode(points, lengths, voxel_size=voxel_size) 34 | points_list.append(points) 35 | lengths_list.append(lengths) 36 | voxel_size *= 2 37 | 38 | # radius search 39 | for i in range(num_stages): 40 | cur_points = points_list[i] 41 | cur_lengths = lengths_list[i] 42 | 43 | neighbors = radius_search_pack_mode( 44 | cur_points, cur_points, cur_lengths, cur_lengths, radius, neighbor_limits[i] 45 | ) 46 | neighbors_list.append(neighbors) 47 | 48 | if i < num_stages - 1: 49 | sub_points = points_list[i + 1] 50 | sub_lengths = lengths_list[i + 1] 51 | 52 | subsampling = radius_search_pack_mode( 53 | sub_points, cur_points, sub_lengths, cur_lengths, radius, neighbor_limits[i] 54 | ) 55 | subsampling_list.append(subsampling) 56 | 57 | upsampling = radius_search_pack_mode( 58 | cur_points, sub_points, cur_lengths, sub_lengths, radius * 2, neighbor_limits[i + 1] 59 | ) 60 | upsampling_list.append(upsampling) 61 | 62 | radius *= 2 63 | 64 | return { 65 | "points": points_list, 66 | "lengths": lengths_list, 67 | "neighbors": neighbors_list, 68 | "subsampling": subsampling_list, 69 | "upsampling": upsampling_list, 70 | } 71 | -------------------------------------------------------------------------------- /vision3d/array_ops/grid_subsample.py: -------------------------------------------------------------------------------- 1 | from vision3d.utils.misc import load_ext 2 | 3 | 4 | ext_module = load_ext("vision3d.ext", ["grid_subsampling"]) 5 | 6 | 7 | def grid_subsample_pack_mode(points, lengths, voxel_size): 8 | """Grid subsampling in pack mode. 9 | 10 | This function is implemented on CPU. 11 | 12 | Args: 13 | points (array): points in pack mode. (N, 3) 14 | lengths (array): number of points in the stacked batch. (B,) 15 | voxel_size (float): voxel size. 16 | 17 | Returns: 18 | s_points (array): sampled points in pack mode (M, 3) 19 | s_lengths (array): numbers of sampled points in the batch. (B,) 20 | """ 21 | s_points, s_lengths = ext_module.grid_subsampling(points, lengths, voxel_size) 22 | return s_points, s_lengths 23 | -------------------------------------------------------------------------------- /vision3d/array_ops/knn.py: -------------------------------------------------------------------------------- 1 | from numpy import ndarray 2 | from scipy.spatial import cKDTree 3 | 4 | 5 | def knn(q_points: ndarray, s_points: ndarray, k=1, return_distance: bool = False): 6 | """Compute the nearest neighbor for the query points in support points. 7 | 8 | Note: 9 | If k=1, the return arrays are squeezed. 10 | """ 11 | s_tree = cKDTree(s_points) 12 | knn_distances, knn_indices = s_tree.query(q_points, k=k) 13 | if return_distance: 14 | return knn_distances, knn_indices 15 | else: 16 | return knn_indices 17 | -------------------------------------------------------------------------------- /vision3d/array_ops/knn_interpolate.py: -------------------------------------------------------------------------------- 1 | from numpy import ndarray 2 | 3 | from .knn import knn 4 | 5 | 6 | def knn_interpolate( 7 | q_points: ndarray, 8 | s_points: ndarray, 9 | s_feats: ndarray, 10 | k: int = 3, 11 | distance_limit: float = 0.1, 12 | eps: float = 1e-10, 13 | inf: float = 1e10, 14 | ) -> ndarray: 15 | distances, indices = knn(q_points, s_points, k=k, return_distance=True) # (N, 3) 16 | if distance_limit is not None: 17 | distances[distances > distance_limit] = inf 18 | weights = 1.0 / (distances + eps) 19 | weights = weights / weights.sum(axis=1, keepdims=True) # (N, 3) 20 | knn_feats = s_feats[indices] # (N, 3, C) 21 | q_feats = (knn_feats * weights[:, :, None]).sum(axis=1) 22 | return q_feats 23 | -------------------------------------------------------------------------------- /vision3d/array_ops/mutual_select.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import ndarray 3 | 4 | from .knn import knn 5 | 6 | 7 | def mutual_select(src_feats: ndarray, tgt_feats: ndarray, mutual: bool = False, bidirectional: bool = False): 8 | """Extract correspondence indices from features. 9 | 10 | Args: 11 | tgt_feats (array): (N, C) 12 | src_feats (array): (M, C) 13 | mutual (bool = False): If True, use mutual matching 14 | bidirectional (bool = False): If True, use directional non-mutual matching, ignored if `mutual` is True. 15 | 16 | Returns: 17 | tgt_corr_indices: (M,) 18 | src_corr_indices: (M,) 19 | """ 20 | src_nn_indices = knn(src_feats, tgt_feats, k=1) 21 | if mutual or bidirectional: 22 | tgt_nn_indices = knn(tgt_feats, src_feats, k=1) 23 | src_indices = np.arange(src_feats.shape[0]) 24 | if mutual: 25 | src_masks = np.equal(tgt_nn_indices[src_nn_indices], src_indices) 26 | src_corr_indices = src_indices[src_masks] 27 | tgt_corr_indices = src_nn_indices[src_corr_indices] 28 | else: 29 | tgt_indices = np.arange(tgt_feats.shape[0]) 30 | src_corr_indices = np.concatenate([src_indices, tgt_nn_indices], axis=0) 31 | tgt_corr_indices = np.concatenate([src_nn_indices, tgt_indices], axis=0) 32 | else: 33 | src_corr_indices = np.arange(src_feats.shape[0]) 34 | tgt_corr_indices = src_nn_indices 35 | return src_corr_indices, tgt_corr_indices 36 | -------------------------------------------------------------------------------- /vision3d/array_ops/procrustes.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import numpy as np 4 | from numpy import ndarray 5 | 6 | from .se3 import get_transform_from_rotation_translation 7 | 8 | 9 | def weighted_procrustes( 10 | src_points: ndarray, 11 | tgt_points: ndarray, 12 | weights: Optional[ndarray] = None, 13 | weight_threshold: float = 0.0, 14 | eps: float = 1e-6, 15 | ) -> ndarray: 16 | """Estimate the alignment transformation with weighted SVD. 17 | 18 | Args: 19 | src_points (array): The correspondence points in the source point cloud in the shape of (N, 3). 20 | tgt_points (array): The correspondence points in the target point cloud in the shape of (N, 3). 21 | weights (array, optional): The weights of the correspondences in the shape of (N,). 22 | weight_threshold (float): The threshold value for the weights. Default: 0. 23 | eps (float): The safe number for division. Default: 1e-6. 24 | 25 | Returns: 26 | A ndarray of the estimated transformation in the shape of (4, 4). 27 | """ 28 | if weights is None: 29 | weights = np.ones(shape=(src_points.shape[0],)) # (N,) 30 | weights[weights < weight_threshold] = 0.0 31 | weights = weights / (weights.sum() + eps) 32 | weights = weights[:, None] # (N, 1) 33 | 34 | src_centroid = np.sum(src_points * weights, axis=0) # (3) 35 | tgt_centroid = np.sum(tgt_points * weights, axis=0) # (3) 36 | 37 | src_points_centered = src_points - src_centroid[None, :] # (N, 3) 38 | tgt_points_centered = tgt_points - tgt_centroid[None, :] # (N, 3) 39 | 40 | H = src_points_centered.T @ (weights * tgt_points_centered) 41 | U, _, Vt = np.linalg.svd(H) # H = USV^T 42 | Ut = U.T 43 | V = Vt.T 44 | eye = np.eye(3) 45 | eye[-1, -1] = np.sign(np.linalg.det(V @ Ut)) 46 | R = V @ eye @ Ut 47 | 48 | t = tgt_centroid - R @ src_centroid 49 | 50 | transform = get_transform_from_rotation_translation(R, t) 51 | 52 | return transform 53 | -------------------------------------------------------------------------------- /vision3d/array_ops/radius_nms.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import numpy as np 4 | from numpy import ndarray 5 | 6 | 7 | def radius_nms(points: ndarray, scores: ndarray, radius: float, num_samples: Optional[int] = None) -> ndarray: 8 | sorted_indices = np.argsort(-scores) 9 | kept_indices = [] 10 | while sorted_indices.size > 0: 11 | i = sorted_indices[0] 12 | kept_indices.append(i) 13 | if num_samples is not None: 14 | if len(kept_indices) >= num_samples: 15 | break 16 | if sorted_indices.size == 1: 17 | break 18 | distances = np.linalg.norm(points[i][None] - points[sorted_indices[1:]], axis=1) 19 | rest_indices = np.nonzero(distances > radius)[0] 20 | if rest_indices.size == 0: 21 | break 22 | sorted_indices = sorted_indices[rest_indices + 1] 23 | return np.asarray(kept_indices) 24 | -------------------------------------------------------------------------------- /vision3d/array_ops/radius_search.py: -------------------------------------------------------------------------------- 1 | from vision3d.utils.misc import load_ext 2 | 3 | 4 | ext_module = load_ext("vision3d.ext", ["radius_neighbors"]) 5 | 6 | 7 | def radius_search_pack_mode(q_points, s_points, q_lengths, s_lengths, radius, neighbor_limit): 8 | """Computes neighbors for a batch of q_points and s_points, apply radius search (in pack mode). 9 | 10 | This function is implemented on CPU. 11 | 12 | Args: 13 | q_points (Tensor): the query points (N, 3) 14 | s_points (Tensor): the support points (M, 3) 15 | q_lengths (Tensor): the list of lengths of batch elements in q_points 16 | s_lengths (Tensor): the list of lengths of batch elements in s_points 17 | radius (float): maximum distance of neighbors 18 | neighbor_limit (int): maximum number of neighbors 19 | 20 | Returns: 21 | neighbors (Tensor): the k nearest neighbors of q_points in s_points (N, k). 22 | Filled with M if there are less than k neighbors. 23 | """ 24 | neighbor_indices = ext_module.radius_neighbors(q_points, s_points, q_lengths, s_lengths, radius) 25 | if neighbor_limit > 0: 26 | neighbor_indices = neighbor_indices[:, :neighbor_limit] 27 | return neighbor_indices 28 | -------------------------------------------------------------------------------- /vision3d/csrc/.clang-format: -------------------------------------------------------------------------------- 1 | BasedOnStyle: Google 2 | BinPackArguments: false 3 | BinPackParameters: false 4 | AlignAfterOpenBracket: AlwaysBreak 5 | ColumnLimit: 120 6 | -------------------------------------------------------------------------------- /vision3d/csrc/common/cuda_helper.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_COMMON_CUDA_HELPER_H_ 2 | #define VISION3D_COMMON_CUDA_HELPER_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include 11 | #include 12 | #include 13 | 14 | namespace vision3d { 15 | 16 | #define THREADS_PER_BLOCK 512 17 | 18 | inline int GET_BLOCKS(const int N, const int num_threads = THREADS_PER_BLOCK) { 19 | int optimal_block_num = (N + num_threads - 1) / num_threads; 20 | int max_block_num = 4096; 21 | return std::min(optimal_block_num, max_block_num); 22 | } 23 | 24 | #define CUDA_1D_KERNEL_LOOP(i, n) \ 25 | for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x) 26 | 27 | } // namespace vision3d 28 | 29 | #endif 30 | -------------------------------------------------------------------------------- /vision3d/csrc/common/numpy_helper.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_COMMON_NUMPY_HELPER_H_ 2 | #define VISION3D_COMMON_NUMPY_HELPER_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #endif -------------------------------------------------------------------------------- /vision3d/csrc/common/torch_helper.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_COMMON_TORCH_HELPER_H_ 2 | #define VISION3D_COMMON_TORCH_HELPER_H_ 3 | 4 | #include 5 | #include 6 | 7 | namespace vision3d { 8 | 9 | #define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor") 10 | 11 | #define CHECK_CPU(x) TORCH_CHECK(!x.device().is_cuda(), #x " must be a CPU tensor") 12 | 13 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 14 | 15 | #define CHECK_CUDA_AND_CONTIGUOUS(x) \ 16 | CHECK_CUDA(x); \ 17 | CHECK_CONTIGUOUS(x) 18 | 19 | #define CHECK_SCALAR_TYPE_INT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, #x " must be an int tensor") 20 | 21 | #define CHECK_SCALAR_TYPE_LONG(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Long, #x " must be an long tensor") 22 | 23 | #define CHECK_SCALAR_TYPE_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, #x " must be a float tensor") 24 | 25 | } // namespace vision3d 26 | 27 | #endif 28 | -------------------------------------------------------------------------------- /vision3d/csrc/cpu/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // Modified from https://github.com/HuguesTHOMAS/KPConv-PyTorch 2 | #include "cloud.h" 3 | 4 | namespace vision3d { 5 | 6 | Point max_point(const std::vector& points) { 7 | Point maxP(points[0]); 8 | 9 | for (const auto& p : points) { 10 | if (p.x > maxP.x) { 11 | maxP.x = p.x; 12 | } 13 | if (p.y > maxP.y) { 14 | maxP.y = p.y; 15 | } 16 | if (p.z > maxP.z) { 17 | maxP.z = p.z; 18 | } 19 | } 20 | 21 | return maxP; 22 | } 23 | 24 | Point min_point(const std::vector& points) { 25 | Point minP(points[0]); 26 | 27 | for (const auto& p : points) { 28 | if (p.x < minP.x) { 29 | minP.x = p.x; 30 | } 31 | if (p.y < minP.y) { 32 | minP.y = p.y; 33 | } 34 | if (p.z < minP.z) { 35 | minP.z = p.z; 36 | } 37 | } 38 | 39 | return minP; 40 | } 41 | 42 | } // namespace vision3d 43 | -------------------------------------------------------------------------------- /vision3d/csrc/cpu/deformation_graph/deformation_graph.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CPU_DEFORMATION_GRAPH_H_ 2 | #define VISION3D_CPU_DEFORMATION_GRAPH_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | #include "../cloud/cloud.h" 13 | #include "../voxel/voxel.h" 14 | #include "numpy_helper.h" 15 | #include "torch_helper.h" 16 | 17 | namespace vision3d { 18 | 19 | const int HASH_BASE = 19997; 20 | 21 | std::tuple< 22 | py::array_t, 23 | py::array_t, 24 | py::array_t, 25 | py::array_t, 26 | py::array_t, 27 | py::array_t> 28 | build_deformation_graph_from_point_cloud( 29 | const py::array_t& vertices, 30 | const py::array_t& node_indices, 31 | int num_neighbors, 32 | int num_anchors, 33 | float max_distance, 34 | float node_coverage); 35 | 36 | } // namespace vision3d 37 | 38 | #endif 39 | -------------------------------------------------------------------------------- /vision3d/csrc/cpu/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CPU_GRID_SUBSAMPLING_H_ 2 | #define VISION3D_CPU_GRID_SUBSAMPLING_H_ 3 | 4 | #include 5 | #include 6 | 7 | #include "../cloud/cloud.h" 8 | #include "numpy_helper.h" 9 | #include "torch_helper.h" 10 | 11 | namespace vision3d { 12 | 13 | class SampledData { 14 | public: 15 | int count; 16 | Point point; 17 | 18 | SampledData() { 19 | count = 0; 20 | point = Point(); 21 | } 22 | 23 | void update(const Point& p) { 24 | count += 1; 25 | point += p; 26 | } 27 | }; 28 | 29 | std::tuple, py::array_t> grid_subsampling( 30 | const py::array_t& points, const py::array_t& lengths, float voxel_size); 31 | 32 | } // namespace vision3d 33 | 34 | #endif -------------------------------------------------------------------------------- /vision3d/csrc/cpu/node_sampling/node_sampling.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CPU_NODE_SAMPLING_H_ 2 | #define VISION3D_CPU_NODE_SAMPLING_H_ 3 | 4 | #include 5 | #include 6 | 7 | #include "../cloud/cloud.h" 8 | #include "numpy_helper.h" 9 | #include "torch_helper.h" 10 | 11 | namespace vision3d { 12 | 13 | py::array_t sample_nodes_with_fps( 14 | const py::array_t& points, const float min_distance, const int num_samples); 15 | 16 | } // namespace vision3d 17 | 18 | #endif 19 | -------------------------------------------------------------------------------- /vision3d/csrc/cpu/radius_neighbors/radius_neighbors.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CPU_RADIUS_NEIGHBORS_H_ 2 | #define VISION3D_CPU_RADIUS_NEIGHBORS_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include "../cloud/cloud.h" 9 | #include "numpy_helper.h" 10 | #include "torch_helper.h" 11 | 12 | namespace vision3d { 13 | 14 | typedef nanoflann::KDTreeSingleIndexAdaptor, PointCloud, 3> kd_tree_t; 15 | 16 | py::array_t radius_neighbors( 17 | const py::array_t& q_points, 18 | const py::array_t& s_points, 19 | const py::array_t& q_lengths, 20 | const py::array_t& s_lengths, 21 | float radius); 22 | 23 | } // namespace vision3d 24 | 25 | #endif 26 | -------------------------------------------------------------------------------- /vision3d/csrc/cpu/voxel/voxel.cpp: -------------------------------------------------------------------------------- 1 | #include "voxel.h" 2 | 3 | namespace vision3d { 4 | 5 | Voxel min_voxel(const std::vector& voxels) { 6 | Voxel result(INT_MAX, INT_MAX, INT_MAX); 7 | for (const auto& voxel : voxels) { 8 | if (voxel.x < result.x) { 9 | result.x = voxel.x; 10 | } 11 | if (voxel.y < result.y) { 12 | result.y = voxel.y; 13 | } 14 | if (voxel.z < result.z) { 15 | result.z = voxel.z; 16 | } 17 | } 18 | return result; 19 | } 20 | 21 | Voxel max_voxel(const std::vector& voxels) { 22 | Voxel result(INT_MIN, INT_MIN, INT_MIN); 23 | for (const auto& voxel : voxels) { 24 | if (voxel.x > result.x) { 25 | result.x = voxel.x; 26 | } 27 | if (voxel.y > result.y) { 28 | result.y = voxel.y; 29 | } 30 | if (voxel.z > result.z) { 31 | result.z = voxel.z; 32 | } 33 | } 34 | return result; 35 | } 36 | 37 | Voxel voxelize(const Point& point, float voxel_size) { 38 | int x = static_cast(std::floor(point.x / voxel_size)); 39 | int y = static_cast(std::floor(point.y / voxel_size)); 40 | int z = static_cast(std::floor(point.z / voxel_size)); 41 | return Voxel(x, y, z); 42 | } 43 | 44 | } // namespace vision3d 45 | -------------------------------------------------------------------------------- /vision3d/csrc/cpu/voxel/voxel.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CPU_VOXEL_H_ 2 | #define VISION3D_CPU_VOXEL_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include "../cloud/cloud.h" 9 | 10 | namespace vision3d { 11 | 12 | class Voxel { 13 | public: 14 | int x, y, z; 15 | 16 | Voxel() : x(0), y(0), z(0) {} 17 | Voxel(int x_, int y_, int z_) : x(x_), y(y_), z(z_) {} 18 | 19 | Voxel& operator-=(const Voxel& other) { 20 | this->x -= other.x; 21 | this->y -= other.y; 22 | this->z -= other.z; 23 | return *this; 24 | } 25 | 26 | Voxel& operator+=(const Voxel& other) { 27 | this->x += other.x; 28 | this->y += other.y; 29 | this->z += other.z; 30 | return *this; 31 | } 32 | }; 33 | 34 | inline Voxel operator+(const Voxel& A, const Voxel& B) { return Voxel(A.x + B.x, A.y + B.y, A.z + B.z); } 35 | 36 | inline Voxel operator-(const Voxel& A, const Voxel& B) { return Voxel(A.x - B.x, A.y - B.y, A.z - B.z); } 37 | 38 | Voxel min_voxel(const std::vector& voxels); 39 | 40 | Voxel max_voxel(const std::vector& voxels); 41 | 42 | Voxel voxelize(const Point& point, float voxel_size); 43 | 44 | } // namespace vision3d 45 | 46 | #endif 47 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/ball_query/ball_query.cpp: -------------------------------------------------------------------------------- 1 | #include "ball_query.h" 2 | 3 | namespace vision3d { 4 | 5 | void ball_query( 6 | const at::Tensor& q_points, 7 | const at::Tensor& s_points, 8 | at::Tensor& indices, 9 | const float max_radius, 10 | const int num_samples) { 11 | ball_query_cuda_launcher(q_points, s_points, indices, max_radius, num_samples); 12 | } 13 | 14 | } // namespace vision3d 15 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/ball_query/ball_query.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_BALL_QUERY_H_ 2 | #define VISION3D_CUDA_BALL_QUERY_H_ 3 | 4 | #include "torch_helper.h" 5 | 6 | namespace vision3d { 7 | 8 | void ball_query( 9 | const at::Tensor& q_points, 10 | const at::Tensor& s_points, 11 | at::Tensor& indices, 12 | const float max_radius, 13 | const int num_samples); 14 | 15 | void ball_query_cuda_launcher( 16 | const at::Tensor& q_points, 17 | const at::Tensor& s_points, 18 | at::Tensor& indices, 19 | const float max_radius, 20 | const int num_samples); 21 | 22 | } // namespace vision3d 23 | 24 | #endif 25 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/ball_query/ball_query_cuda.cu: -------------------------------------------------------------------------------- 1 | #include "ball_query_cuda.cuh" 2 | 3 | namespace vision3d { 4 | 5 | void ball_query_cuda_launcher( 6 | const at::Tensor& q_points, 7 | const at::Tensor& s_points, 8 | at::Tensor& indices, 9 | const float max_radius, 10 | const int num_samples) { 11 | CHECK_CUDA_AND_CONTIGUOUS(q_points); 12 | CHECK_CUDA_AND_CONTIGUOUS(s_points); 13 | CHECK_CUDA_AND_CONTIGUOUS(indices); 14 | CHECK_SCALAR_TYPE_LONG(indices); 15 | 16 | at::cuda::CUDAGuard device_guard(q_points.device()); 17 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 18 | 19 | int batch_size = q_points.size(0); 20 | int num_q_points = q_points.size(1); 21 | int num_s_points = s_points.size(1); 22 | 23 | dim3 grid_dim(GET_BLOCKS(num_q_points, THREADS_PER_BLOCK), batch_size); 24 | dim3 block_dim(THREADS_PER_BLOCK); 25 | 26 | AT_DISPATCH_FLOATING_TYPES(q_points.scalar_type(), "ball_query_kernel", [&] { 27 | ball_query_kernel<<>>( 28 | batch_size, 29 | num_q_points, 30 | num_s_points, 31 | num_samples, 32 | max_radius, 33 | q_points.data_ptr(), 34 | s_points.data_ptr(), 35 | indices.data_ptr()); 36 | }); 37 | 38 | AT_CUDA_CHECK(cudaGetLastError()); 39 | } 40 | 41 | } // namespace vision3d 42 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/ball_query/ball_query_cuda.cuh: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_BALL_QUERY_CUDA_CUH_ 2 | #define VISION3D_CUDA_BALL_QUERY_CUDA_CUH_ 3 | 4 | #include "cuda_helper.h" 5 | #include "torch_helper.h" 6 | 7 | namespace vision3d { 8 | 9 | // input: q_points(b, m, 3) s_points(b, n, 3) 10 | // output: idx(b, m, num_samples) 11 | template 12 | __global__ void ball_query_kernel( 13 | int batch_size, 14 | int num_q_points, 15 | int num_s_points, 16 | int num_samples, 17 | float max_radius, 18 | const T* __restrict__ q_points, 19 | const T* __restrict__ s_points, 20 | long* indices) { 21 | int batch_index = blockIdx.y; 22 | if (batch_index >= batch_size) return; 23 | 24 | q_points += batch_index * num_q_points * 3; 25 | s_points += batch_index * num_s_points * 3; 26 | indices += batch_index * num_q_points * num_samples; 27 | 28 | T max_radius2 = static_cast(max_radius) * static_cast(max_radius); 29 | CUDA_1D_KERNEL_LOOP(q_index, num_q_points) { 30 | const T* cur_q_points = q_points + q_index * 3; 31 | long* cur_indices = indices + q_index * num_samples; 32 | 33 | T qx = cur_q_points[0]; 34 | T qy = cur_q_points[1]; 35 | T qz = cur_q_points[2]; 36 | 37 | int cnt = 0; 38 | for (int s_index = 0; s_index < num_s_points; ++s_index) { 39 | T sx = s_points[s_index * 3 + 0]; 40 | T sy = s_points[s_index * 3 + 1]; 41 | T sz = s_points[s_index * 3 + 2]; 42 | T sq_dist = (qx - sx) * (qx - sx) + (qy - sy) * (qy - sy) + (qz - sz) * (qz - sz); 43 | if (sq_dist < max_radius2) { 44 | if (cnt == 0) { 45 | for (int i = 0; i < num_samples; ++i) { 46 | cur_indices[i] = s_index; 47 | } 48 | } 49 | cur_indices[cnt] = s_index; 50 | ++cnt; 51 | if (cnt >= num_samples) break; 52 | } 53 | } 54 | } 55 | } 56 | 57 | } // namespace vision3d 58 | 59 | #endif 60 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/furthest_point_sample/furthest_point_sample.cpp: -------------------------------------------------------------------------------- 1 | #include "furthest_point_sample.h" 2 | 3 | namespace vision3d { 4 | 5 | void furthest_point_sample( 6 | const at::Tensor& points, at::Tensor& distances, at::Tensor& indices, const int num_samples) { 7 | furthest_point_sample_cuda_launcher(points, distances, indices, num_samples); 8 | } 9 | 10 | } // namespace vision3d 11 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/furthest_point_sample/furthest_point_sample.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_FURTHEST_POINT_SAMPLE_H_ 2 | #define VISION3D_CUDA_FURTHEST_POINT_SAMPLE_H_ 3 | 4 | #include "torch_helper.h" 5 | 6 | namespace vision3d { 7 | 8 | void furthest_point_sample(const at::Tensor& points, at::Tensor& distances, at::Tensor& indices, const int num_samples); 9 | 10 | void furthest_point_sample_cuda_launcher( 11 | const at::Tensor& points, at::Tensor& distances, at::Tensor& indices, const int num_samples); 12 | 13 | } // namespace vision3d 14 | 15 | #endif 16 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/gather/gather.cpp: -------------------------------------------------------------------------------- 1 | #include "gather.h" 2 | 3 | namespace vision3d { 4 | 5 | void gather_forward(const at::Tensor& sources, const at::Tensor& indices, at::Tensor& targets) { 6 | gather_forward_cuda_launcher(sources, indices, targets); 7 | } 8 | 9 | void gather_backward( 10 | const at::Tensor& target_grads, const at::Tensor& indices, at::Tensor& source_grads, const int num_sources) { 11 | gather_backward_cuda_launcher(target_grads, indices, source_grads, num_sources); 12 | } 13 | 14 | } // namespace vision3d 15 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/gather/gather.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_GATHER_H_ 2 | #define VISION3D_CUDA_GATHER_H_ 3 | 4 | #include "torch_helper.h" 5 | 6 | namespace vision3d { 7 | 8 | void gather_forward(const at::Tensor& sources, const at::Tensor& indices, at::Tensor& targets); 9 | 10 | void gather_forward_cuda_launcher(const at::Tensor& sources, const at::Tensor& indices, at::Tensor& targets); 11 | 12 | void gather_backward( 13 | const at::Tensor& target_grads, const at::Tensor& indices, at::Tensor& source_grads, const int num_sources); 14 | 15 | void gather_backward_cuda_launcher( 16 | const at::Tensor& target_grads, const at::Tensor& indices, at::Tensor& source_grads, const int num_sources); 17 | 18 | } // namespace vision3d 19 | 20 | #endif 21 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/gather/gather_cuda.cu: -------------------------------------------------------------------------------- 1 | #include "gather_cuda.cuh" 2 | 3 | namespace vision3d { 4 | 5 | void gather_forward_cuda_launcher(const at::Tensor& sources, const at::Tensor& indices, at::Tensor& targets) { 6 | CHECK_CUDA_AND_CONTIGUOUS(sources); 7 | CHECK_CUDA_AND_CONTIGUOUS(indices); 8 | CHECK_CUDA_AND_CONTIGUOUS(targets); 9 | CHECK_SCALAR_TYPE_LONG(indices); 10 | 11 | at::cuda::CUDAGuard device_guard(sources.device()); 12 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 13 | 14 | int batch_size = sources.size(0); 15 | int num_channels = sources.size(1); 16 | int num_sources = sources.size(2); 17 | int num_samples = indices.size(1); 18 | 19 | dim3 grid_dim(GET_BLOCKS(num_samples, THREADS_PER_BLOCK), num_channels, batch_size); 20 | dim3 block_dim(THREADS_PER_BLOCK); 21 | 22 | AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Bool, sources.scalar_type(), "gather_forward_kernel", [&] { 23 | gather_forward_kernel<<>>( 24 | batch_size, 25 | num_channels, 26 | num_sources, 27 | num_samples, 28 | sources.data_ptr(), 29 | indices.data_ptr(), 30 | targets.data_ptr()); 31 | }); 32 | 33 | AT_CUDA_CHECK(cudaGetLastError()); 34 | } 35 | 36 | void gather_backward_cuda_launcher( 37 | const at::Tensor& target_grads, const at::Tensor& indices, at::Tensor& source_grads, const int num_sources) { 38 | CHECK_CUDA_AND_CONTIGUOUS(target_grads); 39 | CHECK_CUDA_AND_CONTIGUOUS(indices); 40 | CHECK_CUDA_AND_CONTIGUOUS(source_grads); 41 | CHECK_SCALAR_TYPE_LONG(indices); 42 | 43 | at::cuda::CUDAGuard device_guard(target_grads.device()); 44 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 45 | 46 | int batch_size = target_grads.size(0); 47 | int num_channels = target_grads.size(1); 48 | int num_samples = target_grads.size(2); 49 | 50 | dim3 grid_dim(GET_BLOCKS(num_samples, THREADS_PER_BLOCK), num_channels, batch_size); 51 | dim3 block_dim(THREADS_PER_BLOCK); 52 | 53 | AT_DISPATCH_FLOATING_TYPES(target_grads.scalar_type(), "gather_backward_kernel", [&] { 54 | gather_backward_kernel<<>>( 55 | batch_size, 56 | num_channels, 57 | num_sources, 58 | num_samples, 59 | target_grads.data_ptr(), 60 | indices.data_ptr(), 61 | source_grads.data_ptr()); 62 | }); 63 | 64 | AT_CUDA_CHECK(cudaGetLastError()); 65 | } 66 | 67 | } // namespace vision3d 68 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/gather/gather_cuda.cuh: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_GATHER_CUDA_CUH_ 2 | #define VISION3D_CUDA_GATHER_CUDA_CUH_ 3 | 4 | #include "cuda_helper.h" 5 | #include "torch_helper.h" 6 | 7 | namespace vision3d { 8 | 9 | // input: points(b, c, n) idx(b, m) 10 | // output: out(b, c, m) 11 | template 12 | __global__ void gather_forward_kernel( 13 | int batch_size, 14 | int num_channels, 15 | int num_sources, 16 | int num_samples, 17 | const T* __restrict__ sources, 18 | const long* __restrict__ indices, 19 | T* targets) { 20 | int batch_index = blockIdx.z; 21 | int channel_index = blockIdx.y; 22 | 23 | if (batch_index >= batch_size || channel_index >= num_channels) return; 24 | 25 | sources += batch_index * num_channels * num_sources + channel_index * num_sources; 26 | indices += batch_index * num_samples; 27 | targets += batch_index * num_channels * num_samples + channel_index * num_samples; 28 | 29 | CUDA_1D_KERNEL_LOOP(t_index, num_samples) { 30 | int s_index = indices[t_index]; 31 | targets[t_index] = sources[s_index]; 32 | } 33 | } 34 | 35 | // input: grad_out(b, c, m) idx(b, m) 36 | // output: grad_points(b, c, n) 37 | template 38 | __global__ void gather_backward_kernel( 39 | int batch_size, 40 | int num_channels, 41 | int num_sources, 42 | int num_samples, 43 | const T* __restrict__ target_grads, 44 | const long* __restrict__ indices, 45 | T* source_grads) { 46 | int batch_index = blockIdx.z; 47 | int channel_index = blockIdx.y; 48 | 49 | if (batch_index >= batch_size || channel_index >= num_channels) return; 50 | 51 | source_grads += batch_index * num_channels * num_sources + channel_index * num_sources; 52 | indices += batch_index * num_samples; 53 | target_grads += batch_index * num_channels * num_samples + channel_index * num_samples; 54 | 55 | CUDA_1D_KERNEL_LOOP(t_index, num_samples) { 56 | int s_index = indices[t_index]; 57 | gpuAtomicAdd(source_grads + s_index, target_grads[t_index]); 58 | } 59 | } 60 | 61 | } // namespace vision3d 62 | 63 | #endif 64 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/group_gather/group_gather.cpp: -------------------------------------------------------------------------------- 1 | #include "group_gather.h" 2 | 3 | namespace vision3d { 4 | 5 | void group_gather_forward(const at::Tensor& sources, const at::Tensor& indices, at::Tensor& targets) { 6 | group_gather_forward_cuda_launcher(sources, indices, targets); 7 | } 8 | 9 | void group_gather_backward( 10 | const at::Tensor& target_grads, const at::Tensor& indices, at::Tensor& source_grads, const int num_sources) { 11 | group_gather_backward_cuda_launcher(target_grads, indices, source_grads, num_sources); 12 | } 13 | 14 | } // namespace vision3d 15 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/group_gather/group_gather.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_GROUP_GATHER_H_ 2 | #define VISION3D_CUDA_GROUP_GATHER_H_ 3 | 4 | #include "torch_helper.h" 5 | 6 | namespace vision3d { 7 | 8 | void group_gather_forward(const at::Tensor& sources, const at::Tensor& indices, at::Tensor& targets); 9 | 10 | void group_gather_forward_cuda_launcher(const at::Tensor& sources, const at::Tensor& indices, at::Tensor& targets); 11 | 12 | void group_gather_backward( 13 | const at::Tensor& target_grads, const at::Tensor& indices, at::Tensor& source_grads, const int num_sources); 14 | 15 | void group_gather_backward_cuda_launcher( 16 | const at::Tensor& target_grads, const at::Tensor& indices, at::Tensor& source_grads, const int num_sources); 17 | 18 | } // namespace vision3d 19 | 20 | #endif 21 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/group_gather/group_gather_cuda.cuh: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_GROUP_GATHER_CUDA_CUH_ 2 | #define VISION3D_CUDA_GROUP_GATHER_CUDA_CUH_ 3 | 4 | #include "cuda_helper.h" 5 | #include "torch_helper.h" 6 | 7 | namespace vision3d { 8 | 9 | // input: sources(b, c, n) idx(b, num_sources, num_samples) 10 | // output: out(b, c, num_sources, num_samples) 11 | template 12 | __global__ void group_gather_forward_kernel( 13 | int batch_size, 14 | int num_channels, 15 | int num_sources, 16 | int num_targets, 17 | int num_samples, 18 | const T* __restrict__ sources, 19 | const long* __restrict__ indices, 20 | T* targets) { 21 | int batch_index = blockIdx.z; 22 | int channel_index = blockIdx.y; 23 | if (batch_index >= batch_size || channel_index >= num_channels) return; 24 | 25 | sources += batch_index * num_channels * num_sources + channel_index * num_sources; 26 | targets += batch_index * num_channels * num_targets * num_samples + channel_index * num_targets * num_samples; 27 | indices += batch_index * num_targets * num_samples; 28 | 29 | CUDA_1D_KERNEL_LOOP(target_index, num_targets * num_samples) { 30 | int source_index = indices[target_index]; 31 | targets[target_index] = sources[source_index]; 32 | } 33 | } 34 | 35 | // input: target_grads(b, c, num_sources, num_samples), idx(b, num_sources, 36 | // num_samples) output: source_grads(b, c, n) 37 | template 38 | __global__ void group_gather_backward_kernel( 39 | int batch_size, 40 | int num_channels, 41 | int num_sources, 42 | int num_targets, 43 | int num_samples, 44 | const T* __restrict__ target_grads, 45 | const long* __restrict__ indices, 46 | T* source_grads) { 47 | int batch_index = blockIdx.z; 48 | int channel_index = blockIdx.y; 49 | if (batch_index >= batch_size || channel_index >= num_channels) return; 50 | 51 | source_grads += batch_index * num_channels * num_sources + channel_index * num_sources; 52 | target_grads += batch_index * num_channels * num_targets * num_samples + channel_index * num_targets * num_samples; 53 | indices += batch_index * num_targets * num_samples; 54 | 55 | CUDA_1D_KERNEL_LOOP(target_index, num_targets * num_samples) { 56 | int source_index = indices[target_index]; 57 | gpuAtomicAdd(source_grads + source_index, target_grads[target_index]); 58 | } 59 | } 60 | 61 | } // namespace vision3d 62 | 63 | #endif 64 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/knn_points/knn_points.cpp: -------------------------------------------------------------------------------- 1 | #include "knn_points.h" 2 | 3 | namespace vision3d { 4 | 5 | void knn_points( 6 | const at::Tensor& q_points, 7 | const at::Tensor& s_points, 8 | at::Tensor& knn_distances, 9 | at::Tensor& knn_indices, 10 | int num_neighbors) { 11 | knn_points_cuda_launcher(q_points, s_points, knn_distances, knn_indices, num_neighbors); 12 | } 13 | 14 | } // namespace vision3d 15 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/knn_points/knn_points.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_KNN_POINTS_H_ 2 | #define VISION3D_CUDA_KNN_POINTS_H_ 3 | 4 | #include 5 | 6 | #include "torch_helper.h" 7 | 8 | namespace vision3d { 9 | 10 | void knn_points( 11 | const at::Tensor& q_points, 12 | const at::Tensor& s_points, 13 | at::Tensor& knn_distances, 14 | at::Tensor& knn_indices, 15 | int num_neighbors); 16 | 17 | void knn_points_cuda_launcher( 18 | const at::Tensor& q_points, 19 | const at::Tensor& s_points, 20 | at::Tensor& knn_distances, 21 | at::Tensor& knn_indices, 22 | int num_neighbors); 23 | 24 | } // namespace vision3d 25 | 26 | #endif 27 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/knn_points/knn_points_cuda.cu: -------------------------------------------------------------------------------- 1 | #include "knn_points_cuda.cuh" 2 | 3 | namespace vision3d { 4 | 5 | void knn_points_cuda_launcher( 6 | const at::Tensor& q_points, 7 | const at::Tensor& s_points, 8 | at::Tensor& knn_distances, 9 | at::Tensor& knn_indices, 10 | int num_neighbors) { 11 | CHECK_CUDA_AND_CONTIGUOUS(q_points); 12 | CHECK_CUDA_AND_CONTIGUOUS(s_points); 13 | CHECK_CUDA_AND_CONTIGUOUS(knn_distances); 14 | CHECK_CUDA_AND_CONTIGUOUS(knn_indices); 15 | CHECK_SCALAR_TYPE_LONG(knn_indices); 16 | 17 | at::cuda::CUDAGuard device_guard(q_points.device()); 18 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 19 | 20 | int batch_size = q_points.size(0); 21 | int num_q_points = q_points.size(1); 22 | int num_s_points = s_points.size(1); 23 | 24 | dim3 grid_dim(GET_BLOCKS(num_q_points, THREADS_PER_BLOCK), batch_size); 25 | dim3 block_dim(THREADS_PER_BLOCK); 26 | 27 | AT_DISPATCH_FLOATING_TYPES(q_points.scalar_type(), "knn_points_kernel", [&] { 28 | knn_points_kernel<<>>( 29 | batch_size, 30 | num_q_points, 31 | num_s_points, 32 | num_neighbors, 33 | q_points.data_ptr(), 34 | s_points.data_ptr(), 35 | knn_distances.data_ptr(), 36 | knn_indices.data_ptr()); 37 | }); 38 | 39 | AT_CUDA_CHECK(cudaGetLastError()); 40 | } 41 | 42 | } // namespace vision3d 43 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/three_interpolate/three_interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include "three_interpolate.h" 2 | 3 | namespace vision3d { 4 | 5 | void three_interpolate_forward( 6 | const at::Tensor& sources, const at::Tensor& indices, const at::Tensor& weights, at::Tensor& targets) { 7 | three_interpolate_forward_cuda_launcher(sources, indices, weights, targets); 8 | } 9 | 10 | void three_interpolate_backward( 11 | const at::Tensor& target_grads, 12 | const at::Tensor& indices, 13 | const at::Tensor& weights, 14 | at::Tensor& source_grads, 15 | const int num_sources) { 16 | three_interpolate_backward_cuda_launcher(target_grads, indices, weights, source_grads, num_sources); 17 | } 18 | 19 | } // namespace vision3d 20 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/three_interpolate/three_interpolate.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_THREE_INTERPOLATE_H_ 2 | #define VISION3D_CUDA_THREE_INTERPOLATE_H_ 3 | 4 | #include "torch_helper.h" 5 | 6 | namespace vision3d { 7 | 8 | void three_interpolate_forward( 9 | const at::Tensor& sources, const at::Tensor& indices, const at::Tensor& weights, at::Tensor& targets); 10 | 11 | void three_interpolate_forward_cuda_launcher( 12 | const at::Tensor& sources, const at::Tensor& indices, const at::Tensor& weights, at::Tensor& targets); 13 | 14 | void three_interpolate_backward( 15 | const at::Tensor& target_grads, 16 | const at::Tensor& indices, 17 | const at::Tensor& weights, 18 | at::Tensor& source_grads, 19 | const int num_sources); 20 | 21 | void three_interpolate_backward_cuda_launcher( 22 | const at::Tensor& target_grads, 23 | const at::Tensor& indices, 24 | const at::Tensor& weights, 25 | at::Tensor& source_grads, 26 | const int num_sources); 27 | 28 | } // namespace vision3d 29 | 30 | #endif 31 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/three_nn/three_nn.cpp: -------------------------------------------------------------------------------- 1 | #include "three_nn.h" 2 | 3 | namespace vision3d { 4 | 5 | void three_nn( 6 | const at::Tensor& q_points, const at::Tensor& s_points, at::Tensor& tnn_distances, at::Tensor& tnn_indices) { 7 | three_nn_cuda_launcher(q_points, s_points, tnn_distances, tnn_indices); 8 | } 9 | 10 | } // namespace vision3d 11 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/three_nn/three_nn.h: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_THREE_NN_H_ 2 | #define VISION3D_CUDA_THREE_NN_H_ 3 | 4 | #include 5 | 6 | #include "torch_helper.h" 7 | 8 | namespace vision3d { 9 | 10 | void three_nn( 11 | const at::Tensor& q_points, const at::Tensor& s_points, at::Tensor& tnn_distances, at::Tensor& tnn_indices); 12 | 13 | void three_nn_cuda_launcher( 14 | const at::Tensor& q_points, const at::Tensor& s_points, at::Tensor& tnn_distances, at::Tensor& tnn_indices); 15 | 16 | } // namespace vision3d 17 | 18 | #endif 19 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/three_nn/three_nn_cuda.cu: -------------------------------------------------------------------------------- 1 | #include "three_nn_cuda.cuh" 2 | 3 | namespace vision3d { 4 | 5 | void three_nn_cuda_launcher( 6 | const at::Tensor& q_points, const at::Tensor& s_points, at::Tensor& tnn_distances, at::Tensor& tnn_indices) { 7 | CHECK_CUDA_AND_CONTIGUOUS(q_points); 8 | CHECK_CUDA_AND_CONTIGUOUS(s_points); 9 | CHECK_CUDA_AND_CONTIGUOUS(tnn_distances); 10 | CHECK_CUDA_AND_CONTIGUOUS(tnn_indices); 11 | CHECK_SCALAR_TYPE_LONG(tnn_indices); 12 | 13 | at::cuda::CUDAGuard device_guard(q_points.device()); 14 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(); 15 | 16 | int batch_size = q_points.size(0); 17 | int num_q_points = q_points.size(1); 18 | int num_s_points = s_points.size(1); 19 | 20 | dim3 grid_dim(GET_BLOCKS(num_q_points, THREADS_PER_BLOCK), batch_size); 21 | dim3 block_dim(THREADS_PER_BLOCK); 22 | 23 | AT_DISPATCH_FLOATING_TYPES(q_points.scalar_type(), "three_nn_kernel", [&] { 24 | three_nn_kernel<<>>( 25 | batch_size, 26 | num_q_points, 27 | num_s_points, 28 | q_points.data_ptr(), 29 | s_points.data_ptr(), 30 | tnn_distances.data_ptr(), 31 | tnn_indices.data_ptr()); 32 | }); 33 | 34 | AT_CUDA_CHECK(cudaGetLastError()); 35 | } 36 | 37 | } // namespace vision3d 38 | -------------------------------------------------------------------------------- /vision3d/csrc/cuda/three_nn/three_nn_cuda.cuh: -------------------------------------------------------------------------------- 1 | #ifndef VISION3D_CUDA_THREE_NN_CUDA_CUH_ 2 | #define VISION3D_CUDA_THREE_NN_CUDA_CUH_ 3 | 4 | #include "cuda_helper.h" 5 | #include "torch_helper.h" 6 | 7 | namespace vision3d { 8 | 9 | // input: q_points(b, n, 3) s_points(b, m, 3) 10 | // output: dist2(b, n, 3), idx(b, n, 3) 11 | template 12 | __global__ void three_nn_kernel( 13 | int batch_size, 14 | int num_q_points, 15 | int num_s_points, 16 | const T* __restrict__ q_points, 17 | const T* __restrict__ s_points, 18 | T* tnn_distances, 19 | long* tnn_indices) { 20 | int batch_index = blockIdx.y; 21 | if (batch_index >= batch_size) return; 22 | 23 | q_points += batch_index * num_q_points * 3; 24 | s_points += batch_index * num_s_points * 3; 25 | tnn_distances += batch_index * num_q_points * 3; 26 | tnn_indices += batch_index * num_q_points * 3; 27 | 28 | CUDA_1D_KERNEL_LOOP(q_index, num_q_points) { 29 | const T* cur_q_points = q_points + q_index * 3; 30 | T* cur_tnn_distances = tnn_distances + q_index * 3; 31 | long* cur_tnn_indices = tnn_indices + q_index * 3; 32 | 33 | T qx = cur_q_points[0]; 34 | T qy = cur_q_points[1]; 35 | T qz = cur_q_points[2]; 36 | 37 | double best1 = 1e40; 38 | double best2 = 1e40; 39 | double best3 = 1e40; 40 | int besti1 = -1; 41 | int besti2 = -1; 42 | int besti3 = -1; 43 | 44 | for (int s_index = 0; s_index < num_s_points; ++s_index) { 45 | T sx = s_points[s_index * 3 + 0]; 46 | T sy = s_points[s_index * 3 + 1]; 47 | T sz = s_points[s_index * 3 + 2]; 48 | T sq_dist = (qx - sx) * (qx - sx) + (qy - sy) * (qy - sy) + (qz - sz) * (qz - sz); 49 | if (sq_dist < best1) { 50 | best3 = best2; 51 | besti3 = besti2; 52 | best2 = best1; 53 | besti2 = besti1; 54 | best1 = sq_dist; 55 | besti1 = s_index; 56 | } else if (sq_dist < best2) { 57 | best3 = best2; 58 | besti3 = besti2; 59 | best2 = sq_dist; 60 | besti2 = s_index; 61 | } else if (sq_dist < best3) { 62 | best3 = sq_dist; 63 | besti3 = s_index; 64 | } 65 | } 66 | 67 | cur_tnn_distances[0] = static_cast(sqrt(best1)); 68 | cur_tnn_distances[1] = static_cast(sqrt(best2)); 69 | cur_tnn_distances[2] = static_cast(sqrt(best3)); 70 | 71 | cur_tnn_indices[0] = besti1; 72 | cur_tnn_indices[1] = besti2; 73 | cur_tnn_indices[2] = besti3; 74 | } 75 | } 76 | 77 | } // namespace vision3d 78 | 79 | #endif 80 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/Cholesky: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_CHOLESKY_MODULE_H 9 | #define EIGEN_CHOLESKY_MODULE_H 10 | 11 | #include "Core" 12 | #include "Jacobi" 13 | 14 | #include "src/Core/util/DisableStupidWarnings.h" 15 | 16 | /** \defgroup Cholesky_Module Cholesky module 17 | * 18 | * 19 | * 20 | * This module provides two variants of the Cholesky decomposition for selfadjoint (hermitian) matrices. 21 | * Those decompositions are also accessible via the following methods: 22 | * - MatrixBase::llt() 23 | * - MatrixBase::ldlt() 24 | * - SelfAdjointView::llt() 25 | * - SelfAdjointView::ldlt() 26 | * 27 | * \code 28 | * #include 29 | * \endcode 30 | */ 31 | 32 | #include "src/Cholesky/LLT.h" 33 | #include "src/Cholesky/LDLT.h" 34 | #ifdef EIGEN_USE_LAPACKE 35 | #ifdef EIGEN_USE_MKL 36 | #include "mkl_lapacke.h" 37 | #else 38 | #include "src/misc/lapacke.h" 39 | #endif 40 | #include "src/Cholesky/LLT_LAPACKE.h" 41 | #endif 42 | 43 | #include "src/Core/util/ReenableStupidWarnings.h" 44 | 45 | #endif // EIGEN_CHOLESKY_MODULE_H 46 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 47 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/CholmodSupport: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_CHOLMODSUPPORT_MODULE_H 9 | #define EIGEN_CHOLMODSUPPORT_MODULE_H 10 | 11 | #include "SparseCore" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | extern "C" { 16 | #include 17 | } 18 | 19 | /** \ingroup Support_modules 20 | * \defgroup CholmodSupport_Module CholmodSupport module 21 | * 22 | * This module provides an interface to the Cholmod library which is part of the suitesparse package. 23 | * It provides the two following main factorization classes: 24 | * - class CholmodSupernodalLLT: a supernodal LLT Cholesky factorization. 25 | * - class CholmodDecomposiiton: a general L(D)LT Cholesky factorization with automatic or explicit runtime selection of the underlying factorization method (supernodal or simplicial). 26 | * 27 | * For the sake of completeness, this module also propose the two following classes: 28 | * - class CholmodSimplicialLLT 29 | * - class CholmodSimplicialLDLT 30 | * Note that these classes does not bring any particular advantage compared to the built-in 31 | * SimplicialLLT and SimplicialLDLT factorization classes. 32 | * 33 | * \code 34 | * #include 35 | * \endcode 36 | * 37 | * In order to use this module, the cholmod headers must be accessible from the include paths, and your binary must be linked to the cholmod library and its dependencies. 38 | * The dependencies depend on how cholmod has been compiled. 39 | * For a cmake based project, you can use our FindCholmod.cmake module to help you in this task. 40 | * 41 | */ 42 | 43 | #include "src/CholmodSupport/CholmodSupport.h" 44 | 45 | #include "src/Core/util/ReenableStupidWarnings.h" 46 | 47 | #endif // EIGEN_CHOLMODSUPPORT_MODULE_H 48 | 49 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/Dense: -------------------------------------------------------------------------------- 1 | #include "Core" 2 | #include "LU" 3 | #include "Cholesky" 4 | #include "QR" 5 | #include "SVD" 6 | #include "Geometry" 7 | #include "Eigenvalues" 8 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/Eigen: -------------------------------------------------------------------------------- 1 | #include "Dense" 2 | #include "Sparse" 3 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/Eigenvalues: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_EIGENVALUES_MODULE_H 9 | #define EIGEN_EIGENVALUES_MODULE_H 10 | 11 | #include "Core" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | #include "Cholesky" 16 | #include "Jacobi" 17 | #include "Householder" 18 | #include "LU" 19 | #include "Geometry" 20 | 21 | /** \defgroup Eigenvalues_Module Eigenvalues module 22 | * 23 | * 24 | * 25 | * This module mainly provides various eigenvalue solvers. 26 | * This module also provides some MatrixBase methods, including: 27 | * - MatrixBase::eigenvalues(), 28 | * - MatrixBase::operatorNorm() 29 | * 30 | * \code 31 | * #include 32 | * \endcode 33 | */ 34 | 35 | #include "src/misc/RealSvd2x2.h" 36 | #include "src/Eigenvalues/Tridiagonalization.h" 37 | #include "src/Eigenvalues/RealSchur.h" 38 | #include "src/Eigenvalues/EigenSolver.h" 39 | #include "src/Eigenvalues/SelfAdjointEigenSolver.h" 40 | #include "src/Eigenvalues/GeneralizedSelfAdjointEigenSolver.h" 41 | #include "src/Eigenvalues/HessenbergDecomposition.h" 42 | #include "src/Eigenvalues/ComplexSchur.h" 43 | #include "src/Eigenvalues/ComplexEigenSolver.h" 44 | #include "src/Eigenvalues/RealQZ.h" 45 | #include "src/Eigenvalues/GeneralizedEigenSolver.h" 46 | #include "src/Eigenvalues/MatrixBaseEigenvalues.h" 47 | #ifdef EIGEN_USE_LAPACKE 48 | #ifdef EIGEN_USE_MKL 49 | #include "mkl_lapacke.h" 50 | #else 51 | #include "src/misc/lapacke.h" 52 | #endif 53 | #include "src/Eigenvalues/RealSchur_LAPACKE.h" 54 | #include "src/Eigenvalues/ComplexSchur_LAPACKE.h" 55 | #include "src/Eigenvalues/SelfAdjointEigenSolver_LAPACKE.h" 56 | #endif 57 | 58 | #include "src/Core/util/ReenableStupidWarnings.h" 59 | 60 | #endif // EIGEN_EIGENVALUES_MODULE_H 61 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 62 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/Geometry: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_GEOMETRY_MODULE_H 9 | #define EIGEN_GEOMETRY_MODULE_H 10 | 11 | #include "Core" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | #include "SVD" 16 | #include "LU" 17 | #include 18 | 19 | /** \defgroup Geometry_Module Geometry module 20 | * 21 | * This module provides support for: 22 | * - fixed-size homogeneous transformations 23 | * - translation, scaling, 2D and 3D rotations 24 | * - \link Quaternion quaternions \endlink 25 | * - cross products (\ref MatrixBase::cross, \ref MatrixBase::cross3) 26 | * - orthognal vector generation (\ref MatrixBase::unitOrthogonal) 27 | * - some linear components: \link ParametrizedLine parametrized-lines \endlink and \link Hyperplane hyperplanes \endlink 28 | * - \link AlignedBox axis aligned bounding boxes \endlink 29 | * - \link umeyama least-square transformation fitting \endlink 30 | * 31 | * \code 32 | * #include 33 | * \endcode 34 | */ 35 | 36 | #include "src/Geometry/OrthoMethods.h" 37 | #include "src/Geometry/EulerAngles.h" 38 | 39 | #include "src/Geometry/Homogeneous.h" 40 | #include "src/Geometry/RotationBase.h" 41 | #include "src/Geometry/Rotation2D.h" 42 | #include "src/Geometry/Quaternion.h" 43 | #include "src/Geometry/AngleAxis.h" 44 | #include "src/Geometry/Transform.h" 45 | #include "src/Geometry/Translation.h" 46 | #include "src/Geometry/Scaling.h" 47 | #include "src/Geometry/Hyperplane.h" 48 | #include "src/Geometry/ParametrizedLine.h" 49 | #include "src/Geometry/AlignedBox.h" 50 | #include "src/Geometry/Umeyama.h" 51 | 52 | // Use the SSE optimized version whenever possible. At the moment the 53 | // SSE version doesn't compile when AVX is enabled 54 | #if defined EIGEN_VECTORIZE_SSE && !defined EIGEN_VECTORIZE_AVX 55 | #include "src/Geometry/arch/Geometry_SSE.h" 56 | #endif 57 | 58 | #include "src/Core/util/ReenableStupidWarnings.h" 59 | 60 | #endif // EIGEN_GEOMETRY_MODULE_H 61 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 62 | 63 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/Householder: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_HOUSEHOLDER_MODULE_H 9 | #define EIGEN_HOUSEHOLDER_MODULE_H 10 | 11 | #include "Core" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | /** \defgroup Householder_Module Householder module 16 | * This module provides Householder transformations. 17 | * 18 | * \code 19 | * #include 20 | * \endcode 21 | */ 22 | 23 | #include "src/Householder/Householder.h" 24 | #include "src/Householder/HouseholderSequence.h" 25 | #include "src/Householder/BlockHouseholder.h" 26 | 27 | #include "src/Core/util/ReenableStupidWarnings.h" 28 | 29 | #endif // EIGEN_HOUSEHOLDER_MODULE_H 30 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 31 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/IterativeLinearSolvers: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_ITERATIVELINEARSOLVERS_MODULE_H 9 | #define EIGEN_ITERATIVELINEARSOLVERS_MODULE_H 10 | 11 | #include "SparseCore" 12 | #include "OrderingMethods" 13 | 14 | #include "src/Core/util/DisableStupidWarnings.h" 15 | 16 | /** 17 | * \defgroup IterativeLinearSolvers_Module IterativeLinearSolvers module 18 | * 19 | * This module currently provides iterative methods to solve problems of the form \c A \c x = \c b, where \c A is a squared matrix, usually very large and sparse. 20 | * Those solvers are accessible via the following classes: 21 | * - ConjugateGradient for selfadjoint (hermitian) matrices, 22 | * - LeastSquaresConjugateGradient for rectangular least-square problems, 23 | * - BiCGSTAB for general square matrices. 24 | * 25 | * These iterative solvers are associated with some preconditioners: 26 | * - IdentityPreconditioner - not really useful 27 | * - DiagonalPreconditioner - also called Jacobi preconditioner, work very well on diagonal dominant matrices. 28 | * - IncompleteLUT - incomplete LU factorization with dual thresholding 29 | * 30 | * Such problems can also be solved using the direct sparse decomposition modules: SparseCholesky, CholmodSupport, UmfPackSupport, SuperLUSupport. 31 | * 32 | \code 33 | #include 34 | \endcode 35 | */ 36 | 37 | #include "src/IterativeLinearSolvers/SolveWithGuess.h" 38 | #include "src/IterativeLinearSolvers/IterativeSolverBase.h" 39 | #include "src/IterativeLinearSolvers/BasicPreconditioners.h" 40 | #include "src/IterativeLinearSolvers/ConjugateGradient.h" 41 | #include "src/IterativeLinearSolvers/LeastSquareConjugateGradient.h" 42 | #include "src/IterativeLinearSolvers/BiCGSTAB.h" 43 | #include "src/IterativeLinearSolvers/IncompleteLUT.h" 44 | #include "src/IterativeLinearSolvers/IncompleteCholesky.h" 45 | 46 | #include "src/Core/util/ReenableStupidWarnings.h" 47 | 48 | #endif // EIGEN_ITERATIVELINEARSOLVERS_MODULE_H 49 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/Jacobi: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_JACOBI_MODULE_H 9 | #define EIGEN_JACOBI_MODULE_H 10 | 11 | #include "Core" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | /** \defgroup Jacobi_Module Jacobi module 16 | * This module provides Jacobi and Givens rotations. 17 | * 18 | * \code 19 | * #include 20 | * \endcode 21 | * 22 | * In addition to listed classes, it defines the two following MatrixBase methods to apply a Jacobi or Givens rotation: 23 | * - MatrixBase::applyOnTheLeft() 24 | * - MatrixBase::applyOnTheRight(). 25 | */ 26 | 27 | #include "src/Jacobi/Jacobi.h" 28 | 29 | #include "src/Core/util/ReenableStupidWarnings.h" 30 | 31 | #endif // EIGEN_JACOBI_MODULE_H 32 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 33 | 34 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/LU: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_LU_MODULE_H 9 | #define EIGEN_LU_MODULE_H 10 | 11 | #include "Core" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | /** \defgroup LU_Module LU module 16 | * This module includes %LU decomposition and related notions such as matrix inversion and determinant. 17 | * This module defines the following MatrixBase methods: 18 | * - MatrixBase::inverse() 19 | * - MatrixBase::determinant() 20 | * 21 | * \code 22 | * #include 23 | * \endcode 24 | */ 25 | 26 | #include "src/misc/Kernel.h" 27 | #include "src/misc/Image.h" 28 | #include "src/LU/FullPivLU.h" 29 | #include "src/LU/PartialPivLU.h" 30 | #ifdef EIGEN_USE_LAPACKE 31 | #ifdef EIGEN_USE_MKL 32 | #include "mkl_lapacke.h" 33 | #else 34 | #include "src/misc/lapacke.h" 35 | #endif 36 | #include "src/LU/PartialPivLU_LAPACKE.h" 37 | #endif 38 | #include "src/LU/Determinant.h" 39 | #include "src/LU/InverseImpl.h" 40 | 41 | // Use the SSE optimized version whenever possible. At the moment the 42 | // SSE version doesn't compile when AVX is enabled 43 | #if defined EIGEN_VECTORIZE_SSE && !defined EIGEN_VECTORIZE_AVX 44 | #include "src/LU/arch/Inverse_SSE.h" 45 | #endif 46 | 47 | #include "src/Core/util/ReenableStupidWarnings.h" 48 | 49 | #endif // EIGEN_LU_MODULE_H 50 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 51 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/MetisSupport: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_METISSUPPORT_MODULE_H 9 | #define EIGEN_METISSUPPORT_MODULE_H 10 | 11 | #include "SparseCore" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | extern "C" { 16 | #include 17 | } 18 | 19 | 20 | /** \ingroup Support_modules 21 | * \defgroup MetisSupport_Module MetisSupport module 22 | * 23 | * \code 24 | * #include 25 | * \endcode 26 | * This module defines an interface to the METIS reordering package (http://glaros.dtc.umn.edu/gkhome/views/metis). 27 | * It can be used just as any other built-in method as explained in \link OrderingMethods_Module here. \endlink 28 | */ 29 | 30 | 31 | #include "src/MetisSupport/MetisSupport.h" 32 | 33 | #include "src/Core/util/ReenableStupidWarnings.h" 34 | 35 | #endif // EIGEN_METISSUPPORT_MODULE_H 36 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/PaStiXSupport: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_PASTIXSUPPORT_MODULE_H 9 | #define EIGEN_PASTIXSUPPORT_MODULE_H 10 | 11 | #include "SparseCore" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | extern "C" { 16 | #include 17 | #include 18 | } 19 | 20 | #ifdef complex 21 | #undef complex 22 | #endif 23 | 24 | /** \ingroup Support_modules 25 | * \defgroup PaStiXSupport_Module PaStiXSupport module 26 | * 27 | * This module provides an interface to the PaSTiX library. 28 | * PaSTiX is a general \b supernodal, \b parallel and \b opensource sparse solver. 29 | * It provides the two following main factorization classes: 30 | * - class PastixLLT : a supernodal, parallel LLt Cholesky factorization. 31 | * - class PastixLDLT: a supernodal, parallel LDLt Cholesky factorization. 32 | * - class PastixLU : a supernodal, parallel LU factorization (optimized for a symmetric pattern). 33 | * 34 | * \code 35 | * #include 36 | * \endcode 37 | * 38 | * In order to use this module, the PaSTiX headers must be accessible from the include paths, and your binary must be linked to the PaSTiX library and its dependencies. 39 | * The dependencies depend on how PaSTiX has been compiled. 40 | * For a cmake based project, you can use our FindPaSTiX.cmake module to help you in this task. 41 | * 42 | */ 43 | 44 | #include "src/PaStiXSupport/PaStiXSupport.h" 45 | 46 | #include "src/Core/util/ReenableStupidWarnings.h" 47 | 48 | #endif // EIGEN_PASTIXSUPPORT_MODULE_H 49 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/PardisoSupport: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_PARDISOSUPPORT_MODULE_H 9 | #define EIGEN_PARDISOSUPPORT_MODULE_H 10 | 11 | #include "SparseCore" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | #include 16 | 17 | /** \ingroup Support_modules 18 | * \defgroup PardisoSupport_Module PardisoSupport module 19 | * 20 | * This module brings support for the Intel(R) MKL PARDISO direct sparse solvers. 21 | * 22 | * \code 23 | * #include 24 | * \endcode 25 | * 26 | * In order to use this module, the MKL headers must be accessible from the include paths, and your binary must be linked to the MKL library and its dependencies. 27 | * See this \ref TopicUsingIntelMKL "page" for more information on MKL-Eigen integration. 28 | * 29 | */ 30 | 31 | #include "src/PardisoSupport/PardisoSupport.h" 32 | 33 | #include "src/Core/util/ReenableStupidWarnings.h" 34 | 35 | #endif // EIGEN_PARDISOSUPPORT_MODULE_H 36 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/QR: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_QR_MODULE_H 9 | #define EIGEN_QR_MODULE_H 10 | 11 | #include "Core" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | #include "Cholesky" 16 | #include "Jacobi" 17 | #include "Householder" 18 | 19 | /** \defgroup QR_Module QR module 20 | * 21 | * 22 | * 23 | * This module provides various QR decompositions 24 | * This module also provides some MatrixBase methods, including: 25 | * - MatrixBase::householderQr() 26 | * - MatrixBase::colPivHouseholderQr() 27 | * - MatrixBase::fullPivHouseholderQr() 28 | * 29 | * \code 30 | * #include 31 | * \endcode 32 | */ 33 | 34 | #include "src/QR/HouseholderQR.h" 35 | #include "src/QR/FullPivHouseholderQR.h" 36 | #include "src/QR/ColPivHouseholderQR.h" 37 | #include "src/QR/CompleteOrthogonalDecomposition.h" 38 | #ifdef EIGEN_USE_LAPACKE 39 | #ifdef EIGEN_USE_MKL 40 | #include "mkl_lapacke.h" 41 | #else 42 | #include "src/misc/lapacke.h" 43 | #endif 44 | #include "src/QR/HouseholderQR_LAPACKE.h" 45 | #include "src/QR/ColPivHouseholderQR_LAPACKE.h" 46 | #endif 47 | 48 | #include "src/Core/util/ReenableStupidWarnings.h" 49 | 50 | #endif // EIGEN_QR_MODULE_H 51 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 52 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/QtAlignedMalloc: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_QTMALLOC_MODULE_H 9 | #define EIGEN_QTMALLOC_MODULE_H 10 | 11 | #include "Core" 12 | 13 | #if (!EIGEN_MALLOC_ALREADY_ALIGNED) 14 | 15 | #include "src/Core/util/DisableStupidWarnings.h" 16 | 17 | void *qMalloc(std::size_t size) 18 | { 19 | return Eigen::internal::aligned_malloc(size); 20 | } 21 | 22 | void qFree(void *ptr) 23 | { 24 | Eigen::internal::aligned_free(ptr); 25 | } 26 | 27 | void *qRealloc(void *ptr, std::size_t size) 28 | { 29 | void* newPtr = Eigen::internal::aligned_malloc(size); 30 | std::memcpy(newPtr, ptr, size); 31 | Eigen::internal::aligned_free(ptr); 32 | return newPtr; 33 | } 34 | 35 | #include "src/Core/util/ReenableStupidWarnings.h" 36 | 37 | #endif 38 | 39 | #endif // EIGEN_QTMALLOC_MODULE_H 40 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 41 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/SPQRSupport: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_SPQRSUPPORT_MODULE_H 9 | #define EIGEN_SPQRSUPPORT_MODULE_H 10 | 11 | #include "SparseCore" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | #include "SuiteSparseQR.hpp" 16 | 17 | /** \ingroup Support_modules 18 | * \defgroup SPQRSupport_Module SuiteSparseQR module 19 | * 20 | * This module provides an interface to the SPQR library, which is part of the suitesparse package. 21 | * 22 | * \code 23 | * #include 24 | * \endcode 25 | * 26 | * In order to use this module, the SPQR headers must be accessible from the include paths, and your binary must be linked to the SPQR library and its dependencies (Cholmod, AMD, COLAMD,...). 27 | * For a cmake based project, you can use our FindSPQR.cmake and FindCholmod.Cmake modules 28 | * 29 | */ 30 | 31 | #include "src/CholmodSupport/CholmodSupport.h" 32 | #include "src/SPQRSupport/SuiteSparseQRSupport.h" 33 | 34 | #endif 35 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/SVD: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_SVD_MODULE_H 9 | #define EIGEN_SVD_MODULE_H 10 | 11 | #include "QR" 12 | #include "Householder" 13 | #include "Jacobi" 14 | 15 | #include "src/Core/util/DisableStupidWarnings.h" 16 | 17 | /** \defgroup SVD_Module SVD module 18 | * 19 | * 20 | * 21 | * This module provides SVD decomposition for matrices (both real and complex). 22 | * Two decomposition algorithms are provided: 23 | * - JacobiSVD implementing two-sided Jacobi iterations is numerically very accurate, fast for small matrices, but very slow for larger ones. 24 | * - BDCSVD implementing a recursive divide & conquer strategy on top of an upper-bidiagonalization which remains fast for large problems. 25 | * These decompositions are accessible via the respective classes and following MatrixBase methods: 26 | * - MatrixBase::jacobiSvd() 27 | * - MatrixBase::bdcSvd() 28 | * 29 | * \code 30 | * #include 31 | * \endcode 32 | */ 33 | 34 | #include "src/misc/RealSvd2x2.h" 35 | #include "src/SVD/UpperBidiagonalization.h" 36 | #include "src/SVD/SVDBase.h" 37 | #include "src/SVD/JacobiSVD.h" 38 | #include "src/SVD/BDCSVD.h" 39 | #if defined(EIGEN_USE_LAPACKE) && !defined(EIGEN_USE_LAPACKE_STRICT) 40 | #ifdef EIGEN_USE_MKL 41 | #include "mkl_lapacke.h" 42 | #else 43 | #include "src/misc/lapacke.h" 44 | #endif 45 | #include "src/SVD/JacobiSVD_LAPACKE.h" 46 | #endif 47 | 48 | #include "src/Core/util/ReenableStupidWarnings.h" 49 | 50 | #endif // EIGEN_SVD_MODULE_H 51 | /* vim: set filetype=cpp et sw=2 ts=2 ai: */ 52 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/Sparse: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_SPARSE_MODULE_H 9 | #define EIGEN_SPARSE_MODULE_H 10 | 11 | /** \defgroup Sparse_Module Sparse meta-module 12 | * 13 | * Meta-module including all related modules: 14 | * - \ref SparseCore_Module 15 | * - \ref OrderingMethods_Module 16 | * - \ref SparseCholesky_Module 17 | * - \ref SparseLU_Module 18 | * - \ref SparseQR_Module 19 | * - \ref IterativeLinearSolvers_Module 20 | * 21 | \code 22 | #include 23 | \endcode 24 | */ 25 | 26 | #include "SparseCore" 27 | #include "OrderingMethods" 28 | #ifndef EIGEN_MPL2_ONLY 29 | #include "SparseCholesky" 30 | #endif 31 | #include "SparseLU" 32 | #include "SparseQR" 33 | #include "IterativeLinearSolvers" 34 | 35 | #endif // EIGEN_SPARSE_MODULE_H 36 | 37 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/SparseCholesky: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2008-2013 Gael Guennebaud 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_SPARSECHOLESKY_MODULE_H 11 | #define EIGEN_SPARSECHOLESKY_MODULE_H 12 | 13 | #include "SparseCore" 14 | #include "OrderingMethods" 15 | 16 | #include "src/Core/util/DisableStupidWarnings.h" 17 | 18 | /** 19 | * \defgroup SparseCholesky_Module SparseCholesky module 20 | * 21 | * This module currently provides two variants of the direct sparse Cholesky decomposition for selfadjoint (hermitian) matrices. 22 | * Those decompositions are accessible via the following classes: 23 | * - SimplicialLLt, 24 | * - SimplicialLDLt 25 | * 26 | * Such problems can also be solved using the ConjugateGradient solver from the IterativeLinearSolvers module. 27 | * 28 | * \code 29 | * #include 30 | * \endcode 31 | */ 32 | 33 | #ifdef EIGEN_MPL2_ONLY 34 | #error The SparseCholesky module has nothing to offer in MPL2 only mode 35 | #endif 36 | 37 | #include "src/SparseCholesky/SimplicialCholesky.h" 38 | 39 | #ifndef EIGEN_MPL2_ONLY 40 | #include "src/SparseCholesky/SimplicialCholesky_impl.h" 41 | #endif 42 | 43 | #include "src/Core/util/ReenableStupidWarnings.h" 44 | 45 | #endif // EIGEN_SPARSECHOLESKY_MODULE_H 46 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/SparseLU: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2012 Désiré Nuentsa-Wakam 5 | // Copyright (C) 2012 Gael Guennebaud 6 | // 7 | // This Source Code Form is subject to the terms of the Mozilla 8 | // Public License v. 2.0. If a copy of the MPL was not distributed 9 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 10 | 11 | #ifndef EIGEN_SPARSELU_MODULE_H 12 | #define EIGEN_SPARSELU_MODULE_H 13 | 14 | #include "SparseCore" 15 | 16 | /** 17 | * \defgroup SparseLU_Module SparseLU module 18 | * This module defines a supernodal factorization of general sparse matrices. 19 | * The code is fully optimized for supernode-panel updates with specialized kernels. 20 | * Please, see the documentation of the SparseLU class for more details. 21 | */ 22 | 23 | // Ordering interface 24 | #include "OrderingMethods" 25 | 26 | #include "src/SparseLU/SparseLU_gemm_kernel.h" 27 | 28 | #include "src/SparseLU/SparseLU_Structs.h" 29 | #include "src/SparseLU/SparseLU_SupernodalMatrix.h" 30 | #include "src/SparseLU/SparseLUImpl.h" 31 | #include "src/SparseCore/SparseColEtree.h" 32 | #include "src/SparseLU/SparseLU_Memory.h" 33 | #include "src/SparseLU/SparseLU_heap_relax_snode.h" 34 | #include "src/SparseLU/SparseLU_relax_snode.h" 35 | #include "src/SparseLU/SparseLU_pivotL.h" 36 | #include "src/SparseLU/SparseLU_panel_dfs.h" 37 | #include "src/SparseLU/SparseLU_kernel_bmod.h" 38 | #include "src/SparseLU/SparseLU_panel_bmod.h" 39 | #include "src/SparseLU/SparseLU_column_dfs.h" 40 | #include "src/SparseLU/SparseLU_column_bmod.h" 41 | #include "src/SparseLU/SparseLU_copy_to_ucol.h" 42 | #include "src/SparseLU/SparseLU_pruneL.h" 43 | #include "src/SparseLU/SparseLU_Utils.h" 44 | #include "src/SparseLU/SparseLU.h" 45 | 46 | #endif // EIGEN_SPARSELU_MODULE_H 47 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/SparseQR: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_SPARSEQR_MODULE_H 9 | #define EIGEN_SPARSEQR_MODULE_H 10 | 11 | #include "SparseCore" 12 | #include "OrderingMethods" 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | /** \defgroup SparseQR_Module SparseQR module 16 | * \brief Provides QR decomposition for sparse matrices 17 | * 18 | * This module provides a simplicial version of the left-looking Sparse QR decomposition. 19 | * The columns of the input matrix should be reordered to limit the fill-in during the 20 | * decomposition. Built-in methods (COLAMD, AMD) or external methods (METIS) can be used to this end. 21 | * See the \link OrderingMethods_Module OrderingMethods\endlink module for the list 22 | * of built-in and external ordering methods. 23 | * 24 | * \code 25 | * #include 26 | * \endcode 27 | * 28 | * 29 | */ 30 | 31 | #include "OrderingMethods" 32 | #include "src/SparseCore/SparseColEtree.h" 33 | #include "src/SparseQR/SparseQR.h" 34 | 35 | #include "src/Core/util/ReenableStupidWarnings.h" 36 | 37 | #endif 38 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/StdDeque: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2009 Gael Guennebaud 5 | // Copyright (C) 2009 Hauke Heibel 6 | // 7 | // This Source Code Form is subject to the terms of the Mozilla 8 | // Public License v. 2.0. If a copy of the MPL was not distributed 9 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 10 | 11 | #ifndef EIGEN_STDDEQUE_MODULE_H 12 | #define EIGEN_STDDEQUE_MODULE_H 13 | 14 | #include "Core" 15 | #include 16 | 17 | #if EIGEN_COMP_MSVC && EIGEN_OS_WIN64 && (EIGEN_MAX_STATIC_ALIGN_BYTES<=16) /* MSVC auto aligns up to 16 bytes in 64 bit builds */ 18 | 19 | #define EIGEN_DEFINE_STL_DEQUE_SPECIALIZATION(...) 20 | 21 | #else 22 | 23 | #include "src/StlSupport/StdDeque.h" 24 | 25 | #endif 26 | 27 | #endif // EIGEN_STDDEQUE_MODULE_H 28 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/StdList: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2009 Hauke Heibel 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_STDLIST_MODULE_H 11 | #define EIGEN_STDLIST_MODULE_H 12 | 13 | #include "Core" 14 | #include 15 | 16 | #if EIGEN_COMP_MSVC && EIGEN_OS_WIN64 && (EIGEN_MAX_STATIC_ALIGN_BYTES<=16) /* MSVC auto aligns up to 16 bytes in 64 bit builds */ 17 | 18 | #define EIGEN_DEFINE_STL_LIST_SPECIALIZATION(...) 19 | 20 | #else 21 | 22 | #include "src/StlSupport/StdList.h" 23 | 24 | #endif 25 | 26 | #endif // EIGEN_STDLIST_MODULE_H 27 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/StdVector: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2009 Gael Guennebaud 5 | // Copyright (C) 2009 Hauke Heibel 6 | // 7 | // This Source Code Form is subject to the terms of the Mozilla 8 | // Public License v. 2.0. If a copy of the MPL was not distributed 9 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 10 | 11 | #ifndef EIGEN_STDVECTOR_MODULE_H 12 | #define EIGEN_STDVECTOR_MODULE_H 13 | 14 | #include "Core" 15 | #include 16 | 17 | #if EIGEN_COMP_MSVC && EIGEN_OS_WIN64 && (EIGEN_MAX_STATIC_ALIGN_BYTES<=16) /* MSVC auto aligns up to 16 bytes in 64 bit builds */ 18 | 19 | #define EIGEN_DEFINE_STL_VECTOR_SPECIALIZATION(...) 20 | 21 | #else 22 | 23 | #include "src/StlSupport/StdVector.h" 24 | 25 | #endif 26 | 27 | #endif // EIGEN_STDVECTOR_MODULE_H 28 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/SuperLUSupport: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_SUPERLUSUPPORT_MODULE_H 9 | #define EIGEN_SUPERLUSUPPORT_MODULE_H 10 | 11 | #include "SparseCore" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | #ifdef EMPTY 16 | #define EIGEN_EMPTY_WAS_ALREADY_DEFINED 17 | #endif 18 | 19 | typedef int int_t; 20 | #include 21 | #include 22 | #include 23 | 24 | // slu_util.h defines a preprocessor token named EMPTY which is really polluting, 25 | // so we remove it in favor of a SUPERLU_EMPTY token. 26 | // If EMPTY was already defined then we don't undef it. 27 | 28 | #if defined(EIGEN_EMPTY_WAS_ALREADY_DEFINED) 29 | # undef EIGEN_EMPTY_WAS_ALREADY_DEFINED 30 | #elif defined(EMPTY) 31 | # undef EMPTY 32 | #endif 33 | 34 | #define SUPERLU_EMPTY (-1) 35 | 36 | namespace Eigen { struct SluMatrix; } 37 | 38 | /** \ingroup Support_modules 39 | * \defgroup SuperLUSupport_Module SuperLUSupport module 40 | * 41 | * This module provides an interface to the SuperLU library. 42 | * It provides the following factorization class: 43 | * - class SuperLU: a supernodal sequential LU factorization. 44 | * - class SuperILU: a supernodal sequential incomplete LU factorization (to be used as a preconditioner for iterative methods). 45 | * 46 | * \warning This wrapper requires at least versions 4.0 of SuperLU. The 3.x versions are not supported. 47 | * 48 | * \warning When including this module, you have to use SUPERLU_EMPTY instead of EMPTY which is no longer defined because it is too polluting. 49 | * 50 | * \code 51 | * #include 52 | * \endcode 53 | * 54 | * In order to use this module, the superlu headers must be accessible from the include paths, and your binary must be linked to the superlu library and its dependencies. 55 | * The dependencies depend on how superlu has been compiled. 56 | * For a cmake based project, you can use our FindSuperLU.cmake module to help you in this task. 57 | * 58 | */ 59 | 60 | #include "src/SuperLUSupport/SuperLUSupport.h" 61 | 62 | #include "src/Core/util/ReenableStupidWarnings.h" 63 | 64 | #endif // EIGEN_SUPERLUSUPPORT_MODULE_H 65 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/UmfPackSupport: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // This Source Code Form is subject to the terms of the Mozilla 5 | // Public License v. 2.0. If a copy of the MPL was not distributed 6 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 7 | 8 | #ifndef EIGEN_UMFPACKSUPPORT_MODULE_H 9 | #define EIGEN_UMFPACKSUPPORT_MODULE_H 10 | 11 | #include "SparseCore" 12 | 13 | #include "src/Core/util/DisableStupidWarnings.h" 14 | 15 | extern "C" { 16 | #include 17 | } 18 | 19 | /** \ingroup Support_modules 20 | * \defgroup UmfPackSupport_Module UmfPackSupport module 21 | * 22 | * This module provides an interface to the UmfPack library which is part of the suitesparse package. 23 | * It provides the following factorization class: 24 | * - class UmfPackLU: a multifrontal sequential LU factorization. 25 | * 26 | * \code 27 | * #include 28 | * \endcode 29 | * 30 | * In order to use this module, the umfpack headers must be accessible from the include paths, and your binary must be linked to the umfpack library and its dependencies. 31 | * The dependencies depend on how umfpack has been compiled. 32 | * For a cmake based project, you can use our FindUmfPack.cmake module to help you in this task. 33 | * 34 | */ 35 | 36 | #include "src/UmfPackSupport/UmfPackSupport.h" 37 | 38 | #include "src/Core/util/ReenableStupidWarnings.h" 39 | 40 | #endif // EIGEN_UMFPACKSUPPORT_MODULE_H 41 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/DiagonalProduct.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2008 Gael Guennebaud 5 | // Copyright (C) 2007-2009 Benoit Jacob 6 | // 7 | // This Source Code Form is subject to the terms of the Mozilla 8 | // Public License v. 2.0. If a copy of the MPL was not distributed 9 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 10 | 11 | #ifndef EIGEN_DIAGONALPRODUCT_H 12 | #define EIGEN_DIAGONALPRODUCT_H 13 | 14 | namespace Eigen { 15 | 16 | /** \returns the diagonal matrix product of \c *this by the diagonal matrix \a diagonal. 17 | */ 18 | template 19 | template 20 | inline const Product 21 | MatrixBase::operator*(const DiagonalBase &a_diagonal) const 22 | { 23 | return Product(derived(),a_diagonal.derived()); 24 | } 25 | 26 | } // end namespace Eigen 27 | 28 | #endif // EIGEN_DIAGONALPRODUCT_H 29 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/SelfCwiseBinaryOp.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2009-2010 Gael Guennebaud 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_SELFCWISEBINARYOP_H 11 | #define EIGEN_SELFCWISEBINARYOP_H 12 | 13 | namespace Eigen { 14 | 15 | // TODO generalize the scalar type of 'other' 16 | 17 | template 18 | EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Derived& DenseBase::operator*=(const Scalar& other) 19 | { 20 | internal::call_assignment(this->derived(), PlainObject::Constant(rows(),cols(),other), internal::mul_assign_op()); 21 | return derived(); 22 | } 23 | 24 | template 25 | EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Derived& ArrayBase::operator+=(const Scalar& other) 26 | { 27 | internal::call_assignment(this->derived(), PlainObject::Constant(rows(),cols(),other), internal::add_assign_op()); 28 | return derived(); 29 | } 30 | 31 | template 32 | EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Derived& ArrayBase::operator-=(const Scalar& other) 33 | { 34 | internal::call_assignment(this->derived(), PlainObject::Constant(rows(),cols(),other), internal::sub_assign_op()); 35 | return derived(); 36 | } 37 | 38 | template 39 | EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Derived& DenseBase::operator/=(const Scalar& other) 40 | { 41 | internal::call_assignment(this->derived(), PlainObject::Constant(rows(),cols(),other), internal::div_assign_op()); 42 | return derived(); 43 | } 44 | 45 | } // end namespace Eigen 46 | 47 | #endif // EIGEN_SELFCWISEBINARYOP_H 48 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/arch/AVX/TypeCasting.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2015 Benoit Steiner 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_TYPE_CASTING_AVX_H 11 | #define EIGEN_TYPE_CASTING_AVX_H 12 | 13 | namespace Eigen { 14 | 15 | namespace internal { 16 | 17 | // For now we use SSE to handle integers, so we can't use AVX instructions to cast 18 | // from int to float 19 | template <> 20 | struct type_casting_traits { 21 | enum { 22 | VectorizedCast = 0, 23 | SrcCoeffRatio = 1, 24 | TgtCoeffRatio = 1 25 | }; 26 | }; 27 | 28 | template <> 29 | struct type_casting_traits { 30 | enum { 31 | VectorizedCast = 0, 32 | SrcCoeffRatio = 1, 33 | TgtCoeffRatio = 1 34 | }; 35 | }; 36 | 37 | 38 | 39 | template<> EIGEN_STRONG_INLINE Packet8i pcast(const Packet8f& a) { 40 | return _mm256_cvtps_epi32(a); 41 | } 42 | 43 | template<> EIGEN_STRONG_INLINE Packet8f pcast(const Packet8i& a) { 44 | return _mm256_cvtepi32_ps(a); 45 | } 46 | 47 | } // end namespace internal 48 | 49 | } // end namespace Eigen 50 | 51 | #endif // EIGEN_TYPE_CASTING_AVX_H 52 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/arch/Default/ConjHelper.h: -------------------------------------------------------------------------------- 1 | 2 | // This file is part of Eigen, a lightweight C++ template library 3 | // for linear algebra. 4 | // 5 | // Copyright (C) 2017 Gael Guennebaud 6 | // 7 | // This Source Code Form is subject to the terms of the Mozilla 8 | // Public License v. 2.0. If a copy of the MPL was not distributed 9 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 10 | 11 | #ifndef EIGEN_ARCH_CONJ_HELPER_H 12 | #define EIGEN_ARCH_CONJ_HELPER_H 13 | 14 | #define EIGEN_MAKE_CONJ_HELPER_CPLX_REAL(PACKET_CPLX, PACKET_REAL) \ 15 | template<> struct conj_helper { \ 16 | EIGEN_STRONG_INLINE PACKET_CPLX pmadd(const PACKET_REAL& x, const PACKET_CPLX& y, const PACKET_CPLX& c) const \ 17 | { return padd(c, pmul(x,y)); } \ 18 | EIGEN_STRONG_INLINE PACKET_CPLX pmul(const PACKET_REAL& x, const PACKET_CPLX& y) const \ 19 | { return PACKET_CPLX(Eigen::internal::pmul(x, y.v)); } \ 20 | }; \ 21 | \ 22 | template<> struct conj_helper { \ 23 | EIGEN_STRONG_INLINE PACKET_CPLX pmadd(const PACKET_CPLX& x, const PACKET_REAL& y, const PACKET_CPLX& c) const \ 24 | { return padd(c, pmul(x,y)); } \ 25 | EIGEN_STRONG_INLINE PACKET_CPLX pmul(const PACKET_CPLX& x, const PACKET_REAL& y) const \ 26 | { return PACKET_CPLX(Eigen::internal::pmul(x.v, y)); } \ 27 | }; 28 | 29 | #endif // EIGEN_ARCH_CONJ_HELPER_H 30 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/arch/Default/Settings.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2008-2010 Gael Guennebaud 5 | // Copyright (C) 2006-2008 Benoit Jacob 6 | // 7 | // This Source Code Form is subject to the terms of the Mozilla 8 | // Public License v. 2.0. If a copy of the MPL was not distributed 9 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 10 | 11 | 12 | /* All the parameters defined in this file can be specialized in the 13 | * architecture specific files, and/or by the user. 14 | * More to come... */ 15 | 16 | #ifndef EIGEN_DEFAULT_SETTINGS_H 17 | #define EIGEN_DEFAULT_SETTINGS_H 18 | 19 | /** Defines the maximal loop size to enable meta unrolling of loops. 20 | * Note that the value here is expressed in Eigen's own notion of "number of FLOPS", 21 | * it does not correspond to the number of iterations or the number of instructions 22 | */ 23 | #ifndef EIGEN_UNROLLING_LIMIT 24 | #define EIGEN_UNROLLING_LIMIT 100 25 | #endif 26 | 27 | /** Defines the threshold between a "small" and a "large" matrix. 28 | * This threshold is mainly used to select the proper product implementation. 29 | */ 30 | #ifndef EIGEN_CACHEFRIENDLY_PRODUCT_THRESHOLD 31 | #define EIGEN_CACHEFRIENDLY_PRODUCT_THRESHOLD 8 32 | #endif 33 | 34 | /** Defines the maximal width of the blocks used in the triangular product and solver 35 | * for vectors (level 2 blas xTRMV and xTRSV). The default is 8. 36 | */ 37 | #ifndef EIGEN_TUNE_TRIANGULAR_PANEL_WIDTH 38 | #define EIGEN_TUNE_TRIANGULAR_PANEL_WIDTH 8 39 | #endif 40 | 41 | 42 | /** Defines the default number of registers available for that architecture. 43 | * Currently it must be 8 or 16. Other values will fail. 44 | */ 45 | #ifndef EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS 46 | #define EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS 8 47 | #endif 48 | 49 | #endif // EIGEN_DEFAULT_SETTINGS_H 50 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/arch/SSE/TypeCasting.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2015 Benoit Steiner 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_TYPE_CASTING_SSE_H 11 | #define EIGEN_TYPE_CASTING_SSE_H 12 | 13 | namespace Eigen { 14 | 15 | namespace internal { 16 | 17 | #ifndef EIGEN_VECTORIZE_AVX 18 | template <> 19 | struct type_casting_traits { 20 | enum { 21 | VectorizedCast = 1, 22 | SrcCoeffRatio = 1, 23 | TgtCoeffRatio = 1 24 | }; 25 | }; 26 | 27 | template <> 28 | struct type_casting_traits { 29 | enum { 30 | VectorizedCast = 1, 31 | SrcCoeffRatio = 1, 32 | TgtCoeffRatio = 1 33 | }; 34 | }; 35 | 36 | template <> 37 | struct type_casting_traits { 38 | enum { 39 | VectorizedCast = 1, 40 | SrcCoeffRatio = 2, 41 | TgtCoeffRatio = 1 42 | }; 43 | }; 44 | 45 | template <> 46 | struct type_casting_traits { 47 | enum { 48 | VectorizedCast = 1, 49 | SrcCoeffRatio = 1, 50 | TgtCoeffRatio = 2 51 | }; 52 | }; 53 | #endif 54 | 55 | template<> EIGEN_STRONG_INLINE Packet4i pcast(const Packet4f& a) { 56 | return _mm_cvttps_epi32(a); 57 | } 58 | 59 | template<> EIGEN_STRONG_INLINE Packet4f pcast(const Packet4i& a) { 60 | return _mm_cvtepi32_ps(a); 61 | } 62 | 63 | template<> EIGEN_STRONG_INLINE Packet4f pcast(const Packet2d& a, const Packet2d& b) { 64 | return _mm_shuffle_ps(_mm_cvtpd_ps(a), _mm_cvtpd_ps(b), (1 << 2) | (1 << 6)); 65 | } 66 | 67 | template<> EIGEN_STRONG_INLINE Packet2d pcast(const Packet4f& a) { 68 | // Simply discard the second half of the input 69 | return _mm_cvtps_pd(a); 70 | } 71 | 72 | 73 | } // end namespace internal 74 | 75 | } // end namespace Eigen 76 | 77 | #endif // EIGEN_TYPE_CASTING_SSE_H 78 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/functors/TernaryFunctors.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2016 Eugene Brevdo 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_TERNARY_FUNCTORS_H 11 | #define EIGEN_TERNARY_FUNCTORS_H 12 | 13 | namespace Eigen { 14 | 15 | namespace internal { 16 | 17 | //---------- associative ternary functors ---------- 18 | 19 | 20 | 21 | } // end namespace internal 22 | 23 | } // end namespace Eigen 24 | 25 | #endif // EIGEN_TERNARY_FUNCTORS_H 26 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/util/NonMPL2.h: -------------------------------------------------------------------------------- 1 | #ifdef EIGEN_MPL2_ONLY 2 | #error Including non-MPL2 code in EIGEN_MPL2_ONLY mode 3 | #endif 4 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/Core/util/ReenableStupidWarnings.h: -------------------------------------------------------------------------------- 1 | #ifdef EIGEN_WARNINGS_DISABLED 2 | #undef EIGEN_WARNINGS_DISABLED 3 | 4 | #ifndef EIGEN_PERMANENTLY_DISABLE_STUPID_WARNINGS 5 | #ifdef _MSC_VER 6 | #pragma warning( pop ) 7 | #elif defined __INTEL_COMPILER 8 | #pragma warning pop 9 | #elif defined __clang__ 10 | #pragma clang diagnostic pop 11 | #elif defined __GNUC__ && (__GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6)) 12 | #pragma GCC diagnostic pop 13 | #endif 14 | 15 | #if defined __NVCC__ 16 | // Don't reenable the diagnostic messages, as it turns out these messages need 17 | // to be disabled at the point of the template instantiation (i.e the user code) 18 | // otherwise they'll be triggered by nvcc. 19 | // #pragma diag_default code_is_unreachable 20 | // #pragma diag_default initialization_not_reachable 21 | // #pragma diag_default 2651 22 | // #pragma diag_default 2653 23 | #endif 24 | 25 | #endif 26 | 27 | #endif // EIGEN_WARNINGS_DISABLED 28 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/SparseCore/MappedSparseMatrix.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2008-2014 Gael Guennebaud 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_MAPPED_SPARSEMATRIX_H 11 | #define EIGEN_MAPPED_SPARSEMATRIX_H 12 | 13 | namespace Eigen { 14 | 15 | /** \deprecated Use Map > 16 | * \class MappedSparseMatrix 17 | * 18 | * \brief Sparse matrix 19 | * 20 | * \param _Scalar the scalar type, i.e. the type of the coefficients 21 | * 22 | * See http://www.netlib.org/linalg/html_templates/node91.html for details on the storage scheme. 23 | * 24 | */ 25 | namespace internal { 26 | template 27 | struct traits > : traits > 28 | {}; 29 | } // end namespace internal 30 | 31 | template 32 | class MappedSparseMatrix 33 | : public Map > 34 | { 35 | typedef Map > Base; 36 | 37 | public: 38 | 39 | typedef typename Base::StorageIndex StorageIndex; 40 | typedef typename Base::Scalar Scalar; 41 | 42 | inline MappedSparseMatrix(Index rows, Index cols, Index nnz, StorageIndex* outerIndexPtr, StorageIndex* innerIndexPtr, Scalar* valuePtr, StorageIndex* innerNonZeroPtr = 0) 43 | : Base(rows, cols, nnz, outerIndexPtr, innerIndexPtr, valuePtr, innerNonZeroPtr) 44 | {} 45 | 46 | /** Empty destructor */ 47 | inline ~MappedSparseMatrix() {} 48 | }; 49 | 50 | namespace internal { 51 | 52 | template 53 | struct evaluator > 54 | : evaluator > > 55 | { 56 | typedef MappedSparseMatrix<_Scalar,_Options,_StorageIndex> XprType; 57 | typedef evaluator > Base; 58 | 59 | evaluator() : Base() {} 60 | explicit evaluator(const XprType &mat) : Base(mat) {} 61 | }; 62 | 63 | } 64 | 65 | } // end namespace Eigen 66 | 67 | #endif // EIGEN_MAPPED_SPARSEMATRIX_H 68 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/SparseCore/SparseFuzzy.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2008-2014 Gael Guennebaud 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_SPARSE_FUZZY_H 11 | #define EIGEN_SPARSE_FUZZY_H 12 | 13 | namespace Eigen { 14 | 15 | template 16 | template 17 | bool SparseMatrixBase::isApprox(const SparseMatrixBase& other, const RealScalar &prec) const 18 | { 19 | const typename internal::nested_eval::type actualA(derived()); 20 | typename internal::conditional::type, 22 | const PlainObject>::type actualB(other.derived()); 23 | 24 | return (actualA - actualB).squaredNorm() <= prec * prec * numext::mini(actualA.squaredNorm(), actualB.squaredNorm()); 25 | } 26 | 27 | } // end namespace Eigen 28 | 29 | #endif // EIGEN_SPARSE_FUZZY_H 30 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/SparseCore/SparseRedux.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2008-2014 Gael Guennebaud 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | #ifndef EIGEN_SPARSEREDUX_H 11 | #define EIGEN_SPARSEREDUX_H 12 | 13 | namespace Eigen { 14 | 15 | template 16 | typename internal::traits::Scalar 17 | SparseMatrixBase::sum() const 18 | { 19 | eigen_assert(rows()>0 && cols()>0 && "you are using a non initialized matrix"); 20 | Scalar res(0); 21 | internal::evaluator thisEval(derived()); 22 | for (Index j=0; j::InnerIterator iter(thisEval,j); iter; ++iter) 24 | res += iter.value(); 25 | return res; 26 | } 27 | 28 | template 29 | typename internal::traits >::Scalar 30 | SparseMatrix<_Scalar,_Options,_Index>::sum() const 31 | { 32 | eigen_assert(rows()>0 && cols()>0 && "you are using a non initialized matrix"); 33 | if(this->isCompressed()) 34 | return Matrix::Map(m_data.valuePtr(), m_data.size()).sum(); 35 | else 36 | return Base::sum(); 37 | } 38 | 39 | template 40 | typename internal::traits >::Scalar 41 | SparseVector<_Scalar,_Options,_Index>::sum() const 42 | { 43 | eigen_assert(rows()>0 && cols()>0 && "you are using a non initialized matrix"); 44 | return Matrix::Map(m_data.valuePtr(), m_data.size()).sum(); 45 | } 46 | 47 | } // end namespace Eigen 48 | 49 | #endif // EIGEN_SPARSEREDUX_H 50 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/SparseLU/SparseLU_Utils.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2012 Désiré Nuentsa-Wakam 5 | // 6 | // This Source Code Form is subject to the terms of the Mozilla 7 | // Public License v. 2.0. If a copy of the MPL was not distributed 8 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 | 10 | 11 | #ifndef EIGEN_SPARSELU_UTILS_H 12 | #define EIGEN_SPARSELU_UTILS_H 13 | 14 | namespace Eigen { 15 | namespace internal { 16 | 17 | /** 18 | * \brief Count Nonzero elements in the factors 19 | */ 20 | template 21 | void SparseLUImpl::countnz(const Index n, Index& nnzL, Index& nnzU, GlobalLU_t& glu) 22 | { 23 | nnzL = 0; 24 | nnzU = (glu.xusub)(n); 25 | Index nsuper = (glu.supno)(n); 26 | Index jlen; 27 | Index i, j, fsupc; 28 | if (n <= 0 ) return; 29 | // For each supernode 30 | for (i = 0; i <= nsuper; i++) 31 | { 32 | fsupc = glu.xsup(i); 33 | jlen = glu.xlsub(fsupc+1) - glu.xlsub(fsupc); 34 | 35 | for (j = fsupc; j < glu.xsup(i+1); j++) 36 | { 37 | nnzL += jlen; 38 | nnzU += j - fsupc + 1; 39 | jlen--; 40 | } 41 | } 42 | } 43 | 44 | /** 45 | * \brief Fix up the data storage lsub for L-subscripts. 46 | * 47 | * It removes the subscripts sets for structural pruning, 48 | * and applies permutation to the remaining subscripts 49 | * 50 | */ 51 | template 52 | void SparseLUImpl::fixupL(const Index n, const IndexVector& perm_r, GlobalLU_t& glu) 53 | { 54 | Index fsupc, i, j, k, jstart; 55 | 56 | StorageIndex nextl = 0; 57 | Index nsuper = (glu.supno)(n); 58 | 59 | // For each supernode 60 | for (i = 0; i <= nsuper; i++) 61 | { 62 | fsupc = glu.xsup(i); 63 | jstart = glu.xlsub(fsupc); 64 | glu.xlsub(fsupc) = nextl; 65 | for (j = jstart; j < glu.xlsub(fsupc + 1); j++) 66 | { 67 | glu.lsub(nextl) = perm_r(glu.lsub(j)); // Now indexed into P*A 68 | nextl++; 69 | } 70 | for (k = fsupc+1; k < glu.xsup(i+1); k++) 71 | glu.xlsub(k) = nextl; // other columns in supernode i 72 | } 73 | 74 | glu.xlsub(n) = nextl; 75 | } 76 | 77 | } // end namespace internal 78 | 79 | } // end namespace Eigen 80 | #endif // EIGEN_SPARSELU_UTILS_H 81 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/misc/RealSvd2x2.h: -------------------------------------------------------------------------------- 1 | // This file is part of Eigen, a lightweight C++ template library 2 | // for linear algebra. 3 | // 4 | // Copyright (C) 2009-2010 Benoit Jacob 5 | // Copyright (C) 2013-2016 Gael Guennebaud 6 | // 7 | // This Source Code Form is subject to the terms of the Mozilla 8 | // Public License v. 2.0. If a copy of the MPL was not distributed 9 | // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 10 | 11 | #ifndef EIGEN_REALSVD2X2_H 12 | #define EIGEN_REALSVD2X2_H 13 | 14 | namespace Eigen { 15 | 16 | namespace internal { 17 | 18 | template 19 | void real_2x2_jacobi_svd(const MatrixType& matrix, Index p, Index q, 20 | JacobiRotation *j_left, 21 | JacobiRotation *j_right) 22 | { 23 | using std::sqrt; 24 | using std::abs; 25 | Matrix m; 26 | m << numext::real(matrix.coeff(p,p)), numext::real(matrix.coeff(p,q)), 27 | numext::real(matrix.coeff(q,p)), numext::real(matrix.coeff(q,q)); 28 | JacobiRotation rot1; 29 | RealScalar t = m.coeff(0,0) + m.coeff(1,1); 30 | RealScalar d = m.coeff(1,0) - m.coeff(0,1); 31 | 32 | if(abs(d) < (std::numeric_limits::min)()) 33 | { 34 | rot1.s() = RealScalar(0); 35 | rot1.c() = RealScalar(1); 36 | } 37 | else 38 | { 39 | // If d!=0, then t/d cannot overflow because the magnitude of the 40 | // entries forming d are not too small compared to the ones forming t. 41 | RealScalar u = t / d; 42 | RealScalar tmp = sqrt(RealScalar(1) + numext::abs2(u)); 43 | rot1.s() = RealScalar(1) / tmp; 44 | rot1.c() = u / tmp; 45 | } 46 | m.applyOnTheLeft(0,1,rot1); 47 | j_right->makeJacobi(m,0,1); 48 | *j_left = rot1 * j_right->transpose(); 49 | } 50 | 51 | } // end namespace internal 52 | 53 | } // end namespace Eigen 54 | 55 | #endif // EIGEN_REALSVD2X2_H 56 | -------------------------------------------------------------------------------- /vision3d/csrc/external/eigen3/Eigen/src/misc/lapacke_mangling.h: -------------------------------------------------------------------------------- 1 | #ifndef LAPACK_HEADER_INCLUDED 2 | #define LAPACK_HEADER_INCLUDED 3 | 4 | #ifndef LAPACK_GLOBAL 5 | #if defined(LAPACK_GLOBAL_PATTERN_LC) || defined(ADD_) 6 | #define LAPACK_GLOBAL(lcname,UCNAME) lcname##_ 7 | #elif defined(LAPACK_GLOBAL_PATTERN_UC) || defined(UPPER) 8 | #define LAPACK_GLOBAL(lcname,UCNAME) UCNAME 9 | #elif defined(LAPACK_GLOBAL_PATTERN_MC) || defined(NOCHANGE) 10 | #define LAPACK_GLOBAL(lcname,UCNAME) lcname 11 | #else 12 | #define LAPACK_GLOBAL(lcname,UCNAME) lcname##_ 13 | #endif 14 | #endif 15 | 16 | #endif 17 | 18 | -------------------------------------------------------------------------------- /vision3d/datasets/README.md: -------------------------------------------------------------------------------- 1 | # Data Preparation 2 | 3 | ## Object Classification 4 | 5 | ### ModelNet40 6 | 7 | Vision3D uses pre-processed ModelNet40 dataset from PointNet++. 8 | Download the dataset from [here](https://shapenet.cs.stanford.edu/media/modelnet40_normal_resampled.zip). 9 | Modify `data_root` in `tools/datasets/modelnet40.py` and run the file for further pre-processing. 10 | The script will generate 5 hdf5 files for training and 2 hdf5 files for testing with the corresponding id json files. 11 | You can also use the dataset from [here](https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip). 12 | 13 | ## Part Semantic Segmentation 14 | 15 | ### ShapeNetPart 16 | 17 | Vision3D uses pre-processed ShapeNetPart dataset with normal from PointNet++. 18 | Download the dataset from [here](https://shapenet.cs.stanford.edu/media/shapenetcore_partanno_segmentation_benchmark_v0_normal.zip). 19 | Modify `data_root` in `tools/datasets/shapenetpart.py` and run the file for further pre-processing. 20 | The script will generate 4 pickle files for train, val, trainval and test, respectively. 21 | 22 | ## Semantic Segmentation 23 | 24 | ### S3DIS 25 | 26 | Vision3D supports both on-the-fly training augmentation and preprocessed data. 27 | Before using S3DIS in Vision3D, you need to generate `.npy` files following the instructions in [pointnet](https://github.com/charlesq34/pointnet). 28 | For preprocessed hdf5 training data, download from [here](https://shapenet.cs.stanford.edu/media/indoor3d_sem_seg_hdf5_data.zip). 29 | Whole-scene evaluation should be used for testing. 30 | Modify `data_root` in `tools/datastes/s3dis.py` and run the file to generate whole-scene testing hdf5 files for fast evaluation. 31 | 32 | ### ScanNet-v2 33 | 34 | ## Point Cloud Registration 35 | 36 | ### 3DMatch 37 | 38 | ### KITTI odometry 39 | 40 | ### ModelNet40 41 | -------------------------------------------------------------------------------- /vision3d/datasets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qinzheng93/vision3d/7c61728eeeac76c9febf4f70c399358ce1a81b25/vision3d/datasets/__init__.py -------------------------------------------------------------------------------- /vision3d/datasets/classification/__init__.py: -------------------------------------------------------------------------------- 1 | from .modelnet40 import ModelNet40Dataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/classification/modelnet40.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | 3 | import h5py 4 | import numpy as np 5 | import torch.utils.data 6 | 7 | 8 | NUM_CLASSES = 40 9 | # fmt: off 10 | CLASS_NAMES = [ 11 | 'airplane', 'bathtub', 'bed', 'bench', 'bookshelf', 'bottle', 'bowl', 'car', 'chair', 'cone', 'cup', 'curtain', 12 | 'desk', 'door', 'dresser', 'flower_pot', 'glass_box', 'guitar', 'keyboard', 'lamp', 'laptop', 'mantel', 13 | 'monitor', 'night_stand', 'person', 'piano', 'plant', 'radio', 'range_hood', 'sink', 'sofa', 'stairs', 'stool', 14 | 'table', 'tent', 'toilet', 'tv_stand', 'vase', 'wardrobe', 'xbox' 15 | ] 16 | # fmt: on 17 | 18 | 19 | class ModelNet40Dataset(torch.utils.data.Dataset): 20 | def __init__(self, dataset_dir, split, transform): 21 | super().__init__() 22 | 23 | assert split in ["train", "test"] 24 | 25 | self.dataset_dir = dataset_dir 26 | self.transform = transform 27 | 28 | with open(osp.join(self.dataset_dir, "{}_files.txt".format(split))) as f: 29 | lines = f.readlines() 30 | data_files = [line.rstrip() for line in lines] 31 | 32 | self.data = self.read_data(data_files) 33 | 34 | def __getitem__(self, index): 35 | data_dict = { 36 | "points": self.data["points"][index].copy(), 37 | "normals": self.data["normals"][index].copy(), 38 | "label": self.data["labels"][index], 39 | } 40 | if self.transform is not None: 41 | data_dict = self.transform(data_dict) 42 | return data_dict 43 | 44 | def __len__(self): 45 | return self.data["points"].shape[0] 46 | 47 | def read_data(self, data_files): 48 | all_points = [] 49 | all_normals = [] 50 | all_labels = [] 51 | for data_file in data_files: 52 | h5file = h5py.File(osp.join(self.dataset_dir, data_file), "r") 53 | points = h5file["data"][:] 54 | normals = h5file["normal"][:] 55 | labels = h5file["label"][:].flatten().astype(np.int) 56 | all_points.append(points) 57 | all_normals.append(normals) 58 | all_labels.append(labels) 59 | all_points = np.concatenate(all_points, axis=0) 60 | all_normals = np.concatenate(all_normals, axis=0) 61 | all_labels = np.concatenate(all_labels, axis=0) 62 | return {"points": all_points, "normals": all_normals, "labels": all_labels} 63 | -------------------------------------------------------------------------------- /vision3d/datasets/completion/__init__.py: -------------------------------------------------------------------------------- 1 | from .modelnet import ModelNetDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/__init__.py: -------------------------------------------------------------------------------- 1 | from .cape import CapePairDataset 2 | from .deepdeform import DeepDeformPairDataset 3 | from .fourdmatch import FourDMatchPairDataset 4 | from .kitti import OdometryKittiPairDataset 5 | from .modelnet import ModelNetPairDataset 6 | from .rgbdscenes import RGBDScenes2D3DHardPairDataset 7 | from .sevenscenes import SevenScenes2D3DHardPairDataset 8 | from .shapenet import ShapeNetPairDataset 9 | from .threedmatch import ThreeDMatchPairDataset, ThreeDMatchRgbPairDataset 10 | from .threedmatch_urr import ThreeDMatchPairURRDataset 11 | from .scannet_urr import ScanNetPairURRDataset 12 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/cape/__init__.py: -------------------------------------------------------------------------------- 1 | from .cape import CapePairDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/deepdeform/__init__.py: -------------------------------------------------------------------------------- 1 | from .deepdeform import DeepDeformPairDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/fourdmatch/__init__.py: -------------------------------------------------------------------------------- 1 | from .fourdmatch import FourDMatchPairDataset 2 | from .utils import ( 3 | compute_nonrigid_inlier_ratio, 4 | compute_nonrigid_feature_matching_recall, 5 | ) 6 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/kitti/__init__.py: -------------------------------------------------------------------------------- 1 | from .odometry_kitti import OdometryKittiPairDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/modelnet/__init__.py: -------------------------------------------------------------------------------- 1 | from .modelnet import ModelNetPairDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/redwood/__init__.py: -------------------------------------------------------------------------------- 1 | from .redwood import RedwoodPairDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/rgbdscenes/__init__.py: -------------------------------------------------------------------------------- 1 | from .rgbdscenes import RGBDScenes2D3DHardPairDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/scannet_urr/__init__.py: -------------------------------------------------------------------------------- 1 | from .scannet_urr import ScanNetPairURRDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/sevenscenes/__init__.py: -------------------------------------------------------------------------------- 1 | from .sevenscenes_hard import SevenScenes2D3DHardPairDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/shapenet/__init__.py: -------------------------------------------------------------------------------- 1 | from .shapenet import ShapeNetPairDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/threedmatch/__init__.py: -------------------------------------------------------------------------------- 1 | from .threedmatch import ThreeDMatchPairDataset 2 | from .threedmatch_rgb import ThreeDMatchRgbPairDataset 3 | from .threedmatch_utils import ( 4 | calibrate_ground_truth, 5 | compute_transform_error, 6 | evaluate_registration_one_scene, 7 | get_gt_logs_and_infos, 8 | get_num_fragments, 9 | get_scene_abbr, 10 | read_info_file, 11 | read_log_file, 12 | read_pose_file, 13 | write_log_file, 14 | ) 15 | -------------------------------------------------------------------------------- /vision3d/datasets/registration/threedmatch_urr/__init__.py: -------------------------------------------------------------------------------- 1 | from .threedmatch_urr import ThreeDMatchPairURRDataset 2 | -------------------------------------------------------------------------------- /vision3d/datasets/segmentation/__init__.py: -------------------------------------------------------------------------------- 1 | from vision3d.datasets.segmentation.shapenetpart import ShapeNetPartDataset 2 | from vision3d.datasets.segmentation.s3dis import ( 3 | S3DISBlockWiseTrainingDataset, 4 | S3DISBlockWiseTestingDataset, 5 | ) 6 | -------------------------------------------------------------------------------- /vision3d/engine/__init__.py: -------------------------------------------------------------------------------- 1 | from .context_manager import get_context_manager, clear_context_manager 2 | from .batch_tester import BatchTester 3 | from .epoch_based_trainer import EpochBasedTrainer 4 | from .iter_based_trainer import IterBasedTrainer 5 | from .single_tester import SingleTester 6 | from .utils import setup_engine 7 | -------------------------------------------------------------------------------- /vision3d/engine/context_manager.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import warnings 3 | 4 | import torch 5 | from torch import Tensor 6 | 7 | 8 | class ContextManager: 9 | def __init__(self): 10 | self._data_dict = {} 11 | self._scope = [] 12 | 13 | def register(self, name, value, retain_grad=True): 14 | name = osp.join(*self._scope, name) 15 | if isinstance(value, Tensor) and retain_grad: 16 | if value.requires_grad: 17 | value.retain_grad() 18 | self._data_dict[name] = value 19 | 20 | def get(self, name=None, full_name=None): 21 | assert name is not None or full_name is not None 22 | if full_name is None: 23 | full_name = osp.join(*self._scope, name) 24 | if full_name not in self._data_dict: 25 | warnings.warn(f"Key '{full_name}' not found in ContextManager.") 26 | return None 27 | return self._data_dict[full_name] 28 | 29 | def clear(self): 30 | self._data_dict.clear() 31 | self._scope.clear() 32 | 33 | def enter_scope(self, scope_name): 34 | self._scope.append(scope_name) 35 | return self 36 | 37 | def exit_scope(self): 38 | if self._scope: 39 | self._scope.pop() 40 | 41 | @property 42 | def scope(self): 43 | if not self._scope: 44 | return None 45 | return osp.join(*self._scope) 46 | 47 | @property 48 | def data_dict(self): 49 | return self._data_dict 50 | 51 | @property 52 | def data_keys(self): 53 | return [x for x in self._data_dict.keys()] 54 | 55 | def __enter__(self): 56 | return self 57 | 58 | def __exit__(self, exc_type, exc_value, traceback): 59 | self.exit_scope() 60 | 61 | def __repr__(self): 62 | format_string = "ContextManager(\n" 63 | format_string += f" scope='{self.scope}',\n" 64 | format_string += f" names={self.data_keys}\n" 65 | format_string += ")" 66 | return format_string 67 | 68 | 69 | _CONTEXT_MANAGER = ContextManager() 70 | 71 | 72 | def get_context_manager(): 73 | return _CONTEXT_MANAGER 74 | 75 | 76 | def clear_context_manager(): 77 | _CONTEXT_MANAGER.clear() 78 | -------------------------------------------------------------------------------- /vision3d/engine/single_tester.py: -------------------------------------------------------------------------------- 1 | import abc 2 | 3 | import torch 4 | from tqdm import tqdm 5 | 6 | from vision3d.utils.misc import get_log_string 7 | from vision3d.utils.summary_board import SummaryBoard 8 | from vision3d.utils.tensor import move_to_cuda, tensor_to_array 9 | from vision3d.utils.timer import Timer 10 | 11 | from .base_tester import BaseTester 12 | from .context_manager import clear_context_manager 13 | 14 | 15 | class SingleTester(BaseTester, abc.ABC): 16 | def __init__(self, cfg): 17 | super().__init__(cfg) 18 | 19 | def test_epoch(self): 20 | # before epoch 21 | self.before_test_epoch() 22 | # setup watcher 23 | summary_board = SummaryBoard(auto_register=True) 24 | timer = Timer() 25 | # test loop 26 | pbar = tqdm(enumerate(self.test_loader), total=len(self.test_loader)) 27 | timer.tic("data") 28 | for batch_index, data_dict in pbar: 29 | clear_context_manager() 30 | # on start 31 | self.iteration = batch_index + 1 32 | data_dict = move_to_cuda(data_dict) 33 | self.before_test_step(self.iteration, data_dict) 34 | timer.toc("data") 35 | # test step 36 | torch.cuda.synchronize() 37 | timer.tic("model") 38 | output_dict = self.test_step(self.iteration, data_dict) 39 | torch.cuda.synchronize() 40 | timer.toc("model") 41 | # eval step 42 | timer.tic("data") 43 | result_dict = self.eval_step(self.iteration, data_dict, output_dict) 44 | # after step 45 | self.after_test_step(self.iteration, data_dict, output_dict, result_dict) 46 | # logging 47 | result_dict = tensor_to_array(result_dict) 48 | summary_board.update_from_dict(result_dict) 49 | message = self.get_log_string(self.iteration, data_dict, output_dict, result_dict) 50 | pbar.set_description(message + ", " + timer.tostring(keys=["data", "model"], verbose=False)) 51 | torch.cuda.empty_cache() 52 | # summary logging 53 | summary_dict = summary_board.summary() 54 | message = get_log_string(summary_dict, time_dict=timer.summary(keys=["data", "model"])) 55 | self.log(message, level="SUCCESS") 56 | # after epoch 57 | self.after_test_epoch(summary_dict) 58 | -------------------------------------------------------------------------------- /vision3d/engine/utils.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | import numpy as np 4 | import torch 5 | from torch.backends import cudnn as cudnn 6 | 7 | 8 | def setup_engine(seed=None, cudnn_deterministic=True, debug=False): 9 | if seed is not None: 10 | random.seed(seed) 11 | np.random.seed(seed) 12 | torch.manual_seed(seed) 13 | torch.cuda.manual_seed(seed) 14 | torch.cuda.manual_seed_all(seed) 15 | if cudnn_deterministic: 16 | cudnn.benchmark = False 17 | cudnn.deterministic = True 18 | else: 19 | cudnn.benchmark = True 20 | cudnn.deterministic = False 21 | torch.autograd.set_detect_anomaly(debug) 22 | -------------------------------------------------------------------------------- /vision3d/layers/__init__.py: -------------------------------------------------------------------------------- 1 | from .basic_layers import ( 2 | BatchNormPackMode, 3 | DepthwiseConv1d, 4 | DepthwiseConv2d, 5 | GroupNormPackMode, 6 | InstanceNormPackMode, 7 | MonteCarloDropout, 8 | SeparableConv1d, 9 | SeparableConv2d, 10 | build_act_layer, 11 | build_conv_layer, 12 | build_dropout_layer, 13 | build_norm_layer, 14 | build_norm_layer_pack_mode, 15 | check_bias_from_norm_cfg, 16 | convert_monte_carlo_dropout, 17 | ) 18 | from .conv_block import ConvBlock 19 | from .edge_conv import EdgeConv, EdgeConvPackMode 20 | from .embedding import FourierEmbedding, LearnableEmbedding, SinusoidalEmbedding 21 | from .feature_propagate import FeaturePropagate 22 | from .kpconv import InvertedKPResidualBlock, KPConv, KPConvBlock, KPResidualBlock 23 | from .nonrigid_icp import NonRigidICP 24 | from .pointnet import PNConv 25 | from .pointnet2 import GSAConv, SAConv 26 | from .residual_block import BasicConvResBlock, BottleneckConvResBlock 27 | from .sinkhorn import LearnableLogSinkhorn, LogSinkhorn 28 | from .transformer import AttentionLayer, TransformerLayer 29 | from .unary_block import UnaryBlockPackMode 30 | from .weighted_procrustes import WeightedProcrustes 31 | from .xconv import XConv 32 | -------------------------------------------------------------------------------- /vision3d/layers/basic_layers/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import ( 2 | build_act_layer, 3 | build_conv_layer, 4 | build_dropout_layer, 5 | build_norm_layer, 6 | build_norm_layer_pack_mode, 7 | ) 8 | from .depthwise_conv import DepthwiseConv1d, DepthwiseConv2d 9 | from .monte_carlo_dropout import MonteCarloDropout, convert_monte_carlo_dropout 10 | from .norm import BatchNormPackMode, GroupNormPackMode, InstanceNormPackMode 11 | from .separable_conv import SeparableConv1d, SeparableConv2d 12 | from .utils import check_bias_from_norm_cfg 13 | -------------------------------------------------------------------------------- /vision3d/layers/basic_layers/depthwise_conv.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | def _check_depth_multiplier(x): 5 | if not isinstance(x, int) or x <= 0: 6 | raise ValueError(f'"depth_multiplier" ({x}) must be a positive integer.') 7 | 8 | 9 | class DepthwiseConv1d(nn.Conv1d): 10 | def __init__( 11 | self, 12 | in_channels, 13 | kernel_size, 14 | stride=1, 15 | padding=0, 16 | dilation=1, 17 | depth_multiplier=1, 18 | bias=True, 19 | ): 20 | _check_depth_multiplier(depth_multiplier) 21 | out_channels = in_channels * depth_multiplier 22 | super().__init__( 23 | in_channels, 24 | out_channels, 25 | kernel_size=kernel_size, 26 | stride=stride, 27 | padding=padding, 28 | dilation=dilation, 29 | groups=in_channels, 30 | bias=bias, 31 | ) 32 | 33 | 34 | class DepthwiseConv2d(nn.Conv2d): 35 | def __init__( 36 | self, 37 | in_channels, 38 | kernel_size, 39 | stride=1, 40 | padding=0, 41 | dilation=1, 42 | depth_multiplier=1, 43 | bias=True, 44 | ): 45 | _check_depth_multiplier(depth_multiplier) 46 | out_channels = in_channels * depth_multiplier 47 | super().__init__( 48 | in_channels, 49 | out_channels, 50 | kernel_size=kernel_size, 51 | stride=stride, 52 | padding=padding, 53 | dilation=dilation, 54 | groups=in_channels, 55 | bias=bias, 56 | ) 57 | -------------------------------------------------------------------------------- /vision3d/layers/basic_layers/monte_carlo_dropout.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class MonteCarloDropout(nn.Module): 7 | def __init__(self, p=0.1): 8 | super().__init__() 9 | self.p = p 10 | 11 | def forward(self, x): 12 | out = F.dropout(x, p=self.p, training=True) 13 | return out 14 | 15 | 16 | def convert_monte_carlo_dropout(module): 17 | module_output = module 18 | if isinstance(module, torch.nn.Dropout): 19 | module_output = MonteCarloDropout(module.p) 20 | for name, child in module.named_children(): 21 | module_output.add_module(name, convert_monte_carlo_dropout(child)) 22 | del module 23 | return module_output 24 | -------------------------------------------------------------------------------- /vision3d/layers/basic_layers/norm.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch import Tensor 3 | 4 | 5 | class BatchNormPackMode(nn.Module): 6 | def __init__( 7 | self, 8 | num_features: int, 9 | eps: float = 1e-5, 10 | momentum: float = 0.1, 11 | affine: bool = True, 12 | track_running_stats: bool = True, 13 | ): 14 | super().__init__() 15 | self.norm = nn.BatchNorm1d( 16 | num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats 17 | ) 18 | 19 | def forward(self, x: Tensor) -> Tensor: 20 | x = x.transpose(0, 1).unsqueeze(0) # (N, C) -> (B, C, N) 21 | x = self.norm(x) 22 | x = x.squeeze(0).transpose(0, 1) # (B, C, N) -> (N, C) 23 | return x 24 | 25 | def __repr__(self): 26 | return self.__class__.__name__ + "(" + self.norm.extra_repr() + ")" 27 | 28 | 29 | class InstanceNormPackMode(nn.Module): 30 | def __init__( 31 | self, 32 | num_features: int, 33 | eps: float = 1e-5, 34 | momentum: float = 0.1, 35 | affine: bool = False, 36 | track_running_stats: bool = False, 37 | ): 38 | super().__init__() 39 | self.norm = nn.InstanceNorm1d( 40 | num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats 41 | ) 42 | 43 | def forward(self, x: Tensor) -> Tensor: 44 | x = x.transpose(0, 1).unsqueeze(0) # (N, C) -> (B, C, N) 45 | x = self.norm(x) 46 | x = x.squeeze(0).transpose(0, 1) # (B, C, N) -> (N, C) 47 | return x 48 | 49 | def __repr__(self): 50 | return self.__class__.__name__ + "(" + self.norm.extra_repr() + ")" 51 | 52 | 53 | class GroupNormPackMode(nn.Module): 54 | def __init__(self, num_groups: int, num_channels: int, eps: float = 1e-5, affine: bool = True): 55 | super().__init__() 56 | self.norm = nn.GroupNorm(num_groups, num_channels, eps=eps, affine=affine) 57 | 58 | def forward(self, x: Tensor) -> Tensor: 59 | x = x.transpose(0, 1).unsqueeze(0) # (N, C) -> (B, C, N) 60 | x = self.norm(x) 61 | x = x.squeeze(0).transpose(0, 1) # (B, C, N) -> (N, C) 62 | return x 63 | 64 | def __repr__(self): 65 | return self.__class__.__name__ + "(" + self.norm.extra_repr() + ")" 66 | -------------------------------------------------------------------------------- /vision3d/layers/basic_layers/separable_conv.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | def _check_depth_multiplier(x): 5 | if not isinstance(x, int) or x <= 0: 6 | raise ValueError(f'"depth_multiplier" ({x}) must be a positive integer.') 7 | 8 | 9 | class SeparableConv1d(nn.Module): 10 | def __init__( 11 | self, 12 | in_channels, 13 | out_channels, 14 | kernel_size, 15 | stride=1, 16 | padding=0, 17 | dilation=1, 18 | depth_multiplier=1, 19 | bias=True, 20 | padding_mode="zeros", 21 | ): 22 | super().__init__() 23 | _check_depth_multiplier(depth_multiplier) 24 | hidden_dim = in_channels * depth_multiplier 25 | self.dwconv = nn.Conv1d( 26 | in_channels, 27 | hidden_dim, 28 | kernel_size=kernel_size, 29 | stride=stride, 30 | padding=padding, 31 | dilation=dilation, 32 | groups=in_channels, 33 | bias=True, 34 | padding_mode=padding_mode, 35 | ) 36 | self.pwconv = nn.Conv1d(hidden_dim, out_channels, kernel_size=1, bias=bias) 37 | 38 | def forward(self, x): 39 | x = self.dwconv(x) 40 | x = self.pwconv(x) 41 | return x 42 | 43 | 44 | class SeparableConv2d(nn.Module): 45 | def __init__( 46 | self, 47 | in_channels, 48 | out_channels, 49 | kernel_size, 50 | stride=1, 51 | padding=0, 52 | dilation=1, 53 | depth_multiplier=1, 54 | bias=True, 55 | padding_mode="zeros", 56 | ): 57 | super().__init__() 58 | _check_depth_multiplier(depth_multiplier) 59 | dw_channels = in_channels * depth_multiplier 60 | self.dwconv = nn.Conv2d( 61 | in_channels, 62 | dw_channels, 63 | kernel_size=kernel_size, 64 | stride=stride, 65 | padding=padding, 66 | dilation=dilation, 67 | groups=in_channels, 68 | bias=True, 69 | padding_mode=padding_mode, 70 | ) 71 | self.pwconv = nn.Conv2d(dw_channels, out_channels, kernel_size=1, bias=bias) 72 | 73 | def forward(self, x): 74 | x = self.dwconv(x) 75 | x = self.pwconv(x) 76 | return x 77 | -------------------------------------------------------------------------------- /vision3d/layers/basic_layers/utils.py: -------------------------------------------------------------------------------- 1 | from typing import Union 2 | 3 | 4 | def check_bias_from_norm_cfg(norm_cfg: Union[str, dict]) -> bool: 5 | if norm_cfg is None: 6 | return True 7 | if isinstance(norm_cfg, dict): 8 | norm_cfg = norm_cfg["type"] 9 | return not norm_cfg.startswith("BatchNorm") and not norm_cfg.startswith("InstanceNorm") 10 | -------------------------------------------------------------------------------- /vision3d/layers/feature_propagate.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from vision3d.ops import three_interpolate, three_nn 5 | 6 | from .conv_block import ConvBlock 7 | 8 | 9 | class FeaturePropagate(nn.Module): 10 | def __init__(self, input_dim, output_dims, norm_cfg="BatchNorm", act_cfg="ReLU"): 11 | super().__init__() 12 | layers = [] 13 | for output_dim in output_dims: 14 | layers.append( 15 | ConvBlock( 16 | in_channels=input_dim, 17 | out_channels=output_dim, 18 | kernel_size=1, 19 | conv_cfg="Conv1d", 20 | norm_cfg=norm_cfg, 21 | act_cfg=act_cfg, 22 | ) 23 | ) 24 | input_dim = output_dim 25 | self.shared_mlp = nn.Sequential(*layers) 26 | 27 | def forward(self, q_points, s_points, q_feats, s_feats): 28 | """ 29 | Feature Propagation Module. 30 | 31 | Args: 32 | q_points (Tensor): coordinates of the query points (B, 3, Q). 33 | s_points (Tensor): coordinates of the support points (B, 3, S). 34 | q_feats (Tensor): features of the query points (B, Cq, Q). 35 | s_feats (Tensor): features of the support points (B, Cs, S). 36 | 37 | Returns: 38 | q_features (Tensor): output features of the query points (B, CQ + CS, Q). 39 | """ 40 | assert s_points.shape[2] >= 3, f"Too few support points: {s_points.shape}" 41 | 42 | distances, indices = three_nn(q_points, s_points, transposed=True) 43 | weights = torch.div(1.0, distances + 1e-5) 44 | weights = weights / torch.sum(weights, dim=2, keepdim=True) 45 | s_feats = three_interpolate(s_feats, indices, weights) 46 | 47 | if q_feats is not None: 48 | q_feats = torch.cat([q_feats, s_feats], dim=1) 49 | else: 50 | q_feats = s_feats 51 | 52 | q_feats = self.shared_mlp(q_feats) 53 | 54 | return q_feats 55 | -------------------------------------------------------------------------------- /vision3d/layers/kpconv_utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .kernel_points import load_kernels 2 | -------------------------------------------------------------------------------- /vision3d/layers/kpconv_utils/dispositions/k_015_center_3D.ply: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qinzheng93/vision3d/7c61728eeeac76c9febf4f70c399358ce1a81b25/vision3d/layers/kpconv_utils/dispositions/k_015_center_3D.ply -------------------------------------------------------------------------------- /vision3d/layers/pointnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch.nn import functional as F 3 | 4 | from .conv_block import ConvBlock 5 | 6 | 7 | class PNConv(nn.Module): 8 | """PointNet Convolution with a local shared MLP and a global shared MLP. 9 | 10 | Proposed in PointNet. 11 | """ 12 | 13 | def __init__(self, input_dim, local_dims, global_dims, norm_cfg="BatchNorm", act_cfg="ReLU", normalize=False): 14 | super().__init__() 15 | 16 | self.normalize = normalize 17 | 18 | layers = [] 19 | for output_dim in local_dims: 20 | layers.append( 21 | ConvBlock( 22 | in_channels=input_dim, 23 | out_channels=output_dim, 24 | kernel_size=1, 25 | conv_cfg="Conv2d", 26 | norm_cfg=norm_cfg, 27 | act_cfg=act_cfg, 28 | ) 29 | ) 30 | input_dim = output_dim 31 | self.local_mlp = nn.Sequential(*layers) 32 | 33 | layers = [] 34 | for output_dim in global_dims: 35 | layers.append( 36 | ConvBlock( 37 | in_channels=input_dim, 38 | out_channels=output_dim, 39 | kernel_size=1, 40 | conv_cfg="Conv1d", 41 | norm_cfg=norm_cfg, 42 | act_cfg=act_cfg, 43 | ) 44 | ) 45 | input_dim = output_dim 46 | self.global_mlp = nn.Sequential(*layers) 47 | 48 | def forward(self, group_feats): 49 | """PointNet forward. 50 | 51 | Args: 52 | group_feats (Tensor): input features (B, C_in, N, K) 53 | 54 | Returns: 55 | feats (Tensor): output features (B, C_out, N) 56 | """ 57 | group_feats = self.local_mlp(group_feats) 58 | feats = group_feats.max(dim=-1)[0] 59 | feats = self.global_mlp(feats) 60 | if self.normalize: 61 | feats = F.normalize(feats, p=2, dim=1) 62 | return feats 63 | -------------------------------------------------------------------------------- /vision3d/layers/unary_block.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch import Tensor 3 | 4 | from .basic_layers import build_act_layer, build_norm_layer_pack_mode, check_bias_from_norm_cfg 5 | 6 | 7 | class UnaryBlockPackMode(nn.Module): 8 | """Unary block with normalization and activation in pack mode. 9 | 10 | Args: 11 | in_channels (int): dimension input features 12 | out_channels (int): dimension input features 13 | norm_cfg (str|dict|None='GroupNorm'): normalization config 14 | act_cfg (str|dict|None='LeakyReLU'): activation config 15 | """ 16 | 17 | def __init__(self, in_channels, out_channels, norm_cfg="GroupNorm", act_cfg="LeakyReLU"): 18 | super().__init__() 19 | 20 | bias = check_bias_from_norm_cfg(norm_cfg) 21 | self.mlp = nn.Linear(in_channels, out_channels, bias=bias) 22 | 23 | self.norm = build_norm_layer_pack_mode(out_channels, norm_cfg) 24 | self.act = build_act_layer(act_cfg) 25 | 26 | def forward(self, feats: Tensor) -> Tensor: 27 | feats = self.mlp(feats) 28 | feats = self.norm(feats) 29 | feats = self.act(feats) 30 | return feats 31 | -------------------------------------------------------------------------------- /vision3d/layers/weighted_procrustes.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import torch.nn as nn 4 | from torch import Tensor 5 | 6 | from vision3d.ops.weighted_procrustes import weighted_procrustes 7 | 8 | 9 | class WeightedProcrustes(nn.Module): 10 | def __init__(self, weight_threshold: float = 0.0, eps: float = 1e-5): 11 | super().__init__() 12 | self.weight_threshold = weight_threshold 13 | self.eps = eps 14 | 15 | def forward(self, src_points: Tensor, tgt_points: Tensor, weights: Optional[Tensor] = None): 16 | return weighted_procrustes( 17 | src_points, tgt_points, weights=weights, weight_threshold=self.weight_threshold, eps=self.eps 18 | ) 19 | 20 | def extra_repr(self) -> str: 21 | param_strings = [f"weight_threshold={self.weight_threshold:g}", f"eps={self.eps:g}"] 22 | format_string = ", ".join(param_strings) 23 | return format_string 24 | -------------------------------------------------------------------------------- /vision3d/loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .as_rigid_as_possible_loss import AsRigidAsPossibleLoss 2 | from .chamfer_distance import ChamferDistanceLoss 3 | from .circle_loss import CircleLoss 4 | from .focal_loss import SigmoidFocalLoss, SigmoidFocalLossWithLogits 5 | from .hardest_contrastive_loss import HardestContrastiveLoss 6 | from .orthogonal_loss import OrthogonalLoss 7 | from .smooth_ce_loss import SmoothCrossEntropyLoss 8 | from .transformation_loss import TransformationLoss 9 | from .weighted_bce_loss import WeightedBCELoss, WeightedBCEWithLogitsLoss 10 | -------------------------------------------------------------------------------- /vision3d/loss/as_rigid_as_possible_loss.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import torch.nn as nn 4 | from torch import Tensor 5 | 6 | from vision3d.ops import apply_transform, get_rotation_translation_from_transform 7 | 8 | 9 | def as_rigid_as_possible_loss( 10 | nodes: Tensor, transforms: Tensor, edge_indices: Tensor, edge_weights: Optional[Tensor] = None 11 | ) -> Tensor: 12 | """As-rigid-as-possible loss. 13 | 14 | Args: 15 | nodes (Tensor): the nodes in the point clouds in the shape of (V, 3). 16 | transforms (Tensor): the corresponding transformations for the nodes in the shape of (V, 4, 4). 17 | edge_indices (LongTensor): the edge list between the nodes in the shape of (E, 2). 18 | edge_weights (Tensor, optional): the weights of the edges in the shape of (E). 19 | 20 | Returns: 21 | A Tensor of the as-rigid-as-possible-loss. 22 | """ 23 | rotations, translations = get_rotation_translation_from_transform(transforms) # (V, 3, 3), (V, 3) 24 | anc_indices = edge_indices[:, 0] # (E) 25 | ref_indices = edge_indices[:, 1] # (E) 26 | anc_nodes = nodes[anc_indices] # (E, 3) 27 | ref_nodes = nodes[ref_indices] # (E, 3) 28 | anc_transforms = transforms[anc_indices] # (E, 4, 4) 29 | tgt_ref_nodes = translations[ref_indices] + ref_nodes # (E, 3) 30 | warped_ref_nodes = apply_transform(ref_nodes - anc_nodes, anc_transforms) + anc_nodes 31 | loss_values = (warped_ref_nodes - tgt_ref_nodes).pow(2).sum(1) 32 | if edge_weights is not None: 33 | loss = (loss_values * edge_weights).mean() 34 | else: 35 | loss = loss_values.mean() 36 | return loss 37 | 38 | 39 | class AsRigidAsPossibleLoss(nn.Module): 40 | @staticmethod 41 | def forward(nodes: Tensor, transforms: Tensor, edges: Tensor, edge_weights: Optional[Tensor] = None): 42 | loss = as_rigid_as_possible_loss(nodes, transforms, edges, edge_weights=edge_weights) 43 | return loss 44 | -------------------------------------------------------------------------------- /vision3d/loss/orthogonal_loss.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import torch 4 | from torch import Tensor 5 | import torch.nn as nn 6 | from torch.nn import functional as F 7 | 8 | 9 | def orthogonal_loss(inputs, targets: Tensor = None, reduction: str = "mean") -> Tensor: 10 | """Identity loss for rigid transformation. 11 | 12 | If target is None: loss = || R - I ||_F^2, 13 | otherwise: loss = || R^T R* - I ||_F^2 14 | 15 | Args: 16 | inputs (Tensor): input rotations (*, 3, 3) 17 | targets (Tensor=None): target rotations (*, 3, 3) 18 | reduction (str='mean'): reduction method 19 | 20 | Returns: 21 | loss (Tensor): identity loss. 22 | """ 23 | if targets is not None: 24 | inputs = torch.matmul(inputs.transpose(-1, -2), targets) # (*, 3, 3) x (*, 3, 3) -> (*, 3, 3) 25 | inputs = inputs.view(-1, 3, 3) # (*, 3, 3) -> (B, 3, 3) 26 | identity = torch.eye(3).to(inputs).unsqueeze(0).expand_as(inputs) # (*, 3, 3) 27 | loss = F.mse_loss(inputs, identity, reduction=reduction) 28 | return loss 29 | 30 | 31 | class OrthogonalLoss(nn.Module): 32 | def __init__(self, reduction: str = "mean"): 33 | super().__init__() 34 | self.reduction = reduction 35 | 36 | def forward(self, inputs: Tensor, targets: Optional[Tensor] = None) -> Tensor: 37 | return orthogonal_loss(inputs, targets=targets, reduction=self.reduction) 38 | -------------------------------------------------------------------------------- /vision3d/loss/smooth_ce_loss.py: -------------------------------------------------------------------------------- 1 | from torch import Tensor 2 | import torch.nn as nn 3 | from torch.nn import functional as F 4 | 5 | 6 | class SmoothCrossEntropyLoss(nn.Module): 7 | def __init__(self, eps: float = 0.1): 8 | super().__init__() 9 | self.eps = eps 10 | 11 | def forward(self, inputs: Tensor, targets: Tensor) -> Tensor: 12 | """Smooth Cross-entropy Loss. 13 | 14 | Args: 15 | inputs (Tensor): (B, C, *) 16 | targets (LongTensor): (B, *) 17 | 18 | Returns: 19 | loss (Tensor) 20 | """ 21 | batch_size = inputs.shape[0] 22 | num_classes = inputs.shape[1] 23 | inputs = inputs.view(batch_size, num_classes, -1) # (B, C, N) 24 | targets = targets.view(batch_size, -1) # (B, N) 25 | one_hot = F.one_hot(targets, num_classes=num_classes).float().transpose(1, 2) # (B, N, C) -> (B, C, N) 26 | targets = one_hot * (1 - self.eps) + self.eps / inputs.shape[1] 27 | log_inputs = F.log_softmax(inputs, dim=1) 28 | loss = -(targets * log_inputs).sum(dim=1).mean() 29 | return loss 30 | -------------------------------------------------------------------------------- /vision3d/metrics/__init__.py: -------------------------------------------------------------------------------- 1 | from .accuracy import AccuracyMeter 2 | from .iou import IntersectOverUnionMeter 3 | from .part_iou import PartIntersectOverUnionMeter 4 | -------------------------------------------------------------------------------- /vision3d/metrics/accuracy.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import ndarray 3 | 4 | 5 | class AccuracyMeter: 6 | """Accuracy metric for multi-class classification.""" 7 | 8 | def __init__(self, num_classes: int, eps: float = 1e-6): 9 | self.num_classes = num_classes 10 | self.positives = np.zeros(num_classes) 11 | self.counts = np.zeros(num_classes) 12 | self.eps = eps 13 | 14 | def update(self, outputs: ndarray, targets: ndarray): 15 | class_indices = np.arange(self.num_classes)[None, :] # (1, C) 16 | outputs = outputs.flatten()[:, None] # (N, 1) 17 | targets = targets.flatten()[:, None] # (N, 1) 18 | output_mat = outputs == class_indices # (N, C) 19 | target_mat = targets == class_indices # (N, C) 20 | positives = np.sum(output_mat & target_mat, axis=0) # (C,) 21 | counts = np.sum(target_mat, axis=0) # (C,) 22 | self.positives += positives 23 | self.counts += counts 24 | 25 | def accuracy(self, class_id: int): 26 | return self.positives[class_id] / (self.counts[class_id] + self.eps) 27 | 28 | def mean_accuracy(self): 29 | return np.mean(self.positives / (self.counts + self.eps)) 30 | 31 | def overall_accuracy(self): 32 | return np.sum(self.positives) / (np.sum(self.counts) + self.eps) 33 | -------------------------------------------------------------------------------- /vision3d/metrics/iou.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import ndarray 3 | 4 | 5 | class IntersectOverUnionMeter: 6 | """Intersect over Union metric for scene segmentation.""" 7 | 8 | def __init__(self, num_classes: int, eps: float = 1e-6): 9 | self.num_classes = num_classes 10 | self.intersects = np.zeros(num_classes) 11 | self.unions = np.zeros(num_classes) 12 | self.eps = eps 13 | 14 | def update(self, outputs: ndarray, targets: ndarray): 15 | class_indices = np.arange(self.num_classes)[None, :] # (1, C) 16 | outputs = outputs.flatten()[:, None] # (N, 1) 17 | targets = targets.flatten()[:, None] # (N, 1) 18 | output_mat = outputs == class_indices # (N, C) 19 | target_mat = targets == class_indices # (N, C) 20 | intersects = np.sum(output_mat & target_mat, axis=0) # (C,) 21 | unions = np.sum(output_mat | target_mat, axis=0) # (C,) 22 | self.intersects += intersects 23 | self.unions += unions 24 | 25 | def iou(self, class_id: int): 26 | return self.intersects[class_id] / (self.unions[class_id] + self.eps) 27 | 28 | def mean_iou(self): 29 | return np.mean(self.intersects / (self.unions + self.eps)) 30 | 31 | def overall_iou(self): 32 | return np.sum(self.intersects) / (np.sum(self.unions) + self.eps) 33 | -------------------------------------------------------------------------------- /vision3d/metrics/part_iou.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | import numpy as np 3 | from numpy import ndarray 4 | 5 | 6 | class PartIntersectOverUnionMeter: 7 | """Intersect over Union metric for part segmentation.""" 8 | 9 | def __init__(self, num_classes: int, num_parts: int, class_id_to_part_ids: List, eps: float = 1e-6): 10 | self.num_classes = num_classes 11 | self.num_parts = num_parts 12 | self.class_id_to_part_ids = class_id_to_part_ids 13 | self.iou_sums = np.zeros(self.num_classes) 14 | self.counts = np.zeros(self.num_classes) 15 | self.eps = eps 16 | 17 | def update(self, outputs: ndarray, targets: ndarray, class_id: int): 18 | class_indices = np.asarray(self.class_id_to_part_ids[class_id])[None, :] 19 | outputs = outputs.flatten()[:, None] 20 | targets = targets.flatten()[:, None] 21 | output_mat = outputs == class_indices 22 | target_mat = targets == class_indices 23 | intersects = np.sum(output_mat & target_mat, axis=0) 24 | unions = np.sum(output_mat | target_mat, axis=0) 25 | ious = intersects / (unions + self.eps) 26 | ious[unions == 0] = 1.0 27 | iou = np.mean(ious) 28 | self.iou_sums[class_id] += iou 29 | self.counts[class_id] += 1 30 | 31 | def iou(self, class_id: int): 32 | return self.iou_sums[class_id] / (self.counts[class_id] + self.eps) 33 | 34 | def mean_iou(self): 35 | """Mean IoU over all classes.""" 36 | return np.mean(self.iou_sums / (self.counts + self.eps)) 37 | 38 | def overall_iou(self): 39 | """Overall average IoU over all instances.""" 40 | return np.sum(self.iou_sums) / np.sum(self.counts + self.eps) 41 | -------------------------------------------------------------------------------- /vision3d/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qinzheng93/vision3d/7c61728eeeac76c9febf4f70c399358ce1a81b25/vision3d/models/__init__.py -------------------------------------------------------------------------------- /vision3d/models/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | from .kpconv_fpn import KPConvFPN 2 | -------------------------------------------------------------------------------- /vision3d/models/cofinet/__init__.py: -------------------------------------------------------------------------------- 1 | from .node_matching import NodeMatching 2 | from .node_proposal import NodeProposalGenerator 3 | -------------------------------------------------------------------------------- /vision3d/models/cofinet/node_proposal.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from vision3d.ops import random_choice 5 | 6 | 7 | class NodeProposalGenerator(nn.Module): 8 | def __init__(self, num_proposals): 9 | super().__init__() 10 | self.num_proposals = num_proposals 11 | 12 | @torch.no_grad() 13 | def forward(self, gt_src_corr_indices, gt_tgt_corr_indices, gt_corr_overlaps): 14 | """Generate ground truth superpoint (patch) correspondences. 15 | 16 | Randomly select "num_proposals" correspondences whose overlap is above "overlap_threshold". 17 | 18 | Args: 19 | gt_src_corr_indices (LongTensor): ground truth superpoint correspondences (N,) 20 | gt_tgt_corr_indices (LongTensor): ground truth superpoint correspondences (N,) 21 | gt_corr_overlaps (Tensor): ground truth superpoint correspondences overlap (N,) 22 | 23 | Returns: 24 | gt_src_corr_indices (LongTensor): selected superpoints in source point cloud. 25 | gt_tgt_corr_indices (LongTensor): selected superpoints in target point cloud. 26 | gt_corr_overlaps (LongTensor): overlaps of the selected superpoint correspondences. 27 | """ 28 | if gt_corr_overlaps.shape[0] > self.num_proposals: 29 | scores = gt_corr_overlaps / gt_corr_overlaps.sum() 30 | sel_indices = random_choice(gt_corr_overlaps.shape[0], self.num_proposals, replace=False, p=scores) 31 | gt_src_corr_indices = gt_src_corr_indices[sel_indices] 32 | gt_tgt_corr_indices = gt_tgt_corr_indices[sel_indices] 33 | gt_corr_overlaps = gt_corr_overlaps[sel_indices] 34 | 35 | return gt_src_corr_indices, gt_tgt_corr_indices, gt_corr_overlaps 36 | 37 | def extra_repr(self) -> str: 38 | param_strings = [f"num_proposals={self.num_proposals}"] 39 | format_string = ", ".join(param_strings) 40 | return format_string 41 | -------------------------------------------------------------------------------- /vision3d/models/contrib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qinzheng93/vision3d/7c61728eeeac76c9febf4f70c399358ce1a81b25/vision3d/models/contrib/__init__.py -------------------------------------------------------------------------------- /vision3d/models/d3feat/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qinzheng93/vision3d/7c61728eeeac76c9febf4f70c399358ce1a81b25/vision3d/models/d3feat/__init__.py -------------------------------------------------------------------------------- /vision3d/models/fcgf/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /vision3d/models/foldingnet/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /vision3d/models/geotransformer/__init__.py: -------------------------------------------------------------------------------- 1 | from .adaptive_superpoint_matching import AdaptiveSuperPointMatching 2 | from .geometric_structure_embedding import GeometricStructureEmbedding, GeometricStructureEmbeddingV2 3 | from .geometric_transformer import GeometricTransformer 4 | from .geometric_transformer_lite import GeometricTransformerLite 5 | from .local_global_registration import LocalGlobalRegistration 6 | from .local_global_registration_threshold import LocalGlobalRegistrationThreshold 7 | from .local_global_registration_duplicate_removal import LocalGlobalRegistrationWithDuplicateRemoval 8 | from .point_matching import PointMatching 9 | from .spatial_consistency_filtering import SpatialConsistencyFiltering 10 | from .superpoint_matching_global_topk import SuperPointMatchingGlobalTopk 11 | from .superpoint_matching_mutual_topk import SuperPointMatchingMutualTopk 12 | from .superpoint_proposal import SuperPointProposalGenerator 13 | -------------------------------------------------------------------------------- /vision3d/models/geotransformer/spatial_consistency_filtering.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from IPython import embed 4 | from torch import Tensor 5 | 6 | from vision3d.ops import spatial_consistency 7 | 8 | 9 | class SpatialConsistencyFiltering(nn.Module): 10 | def __init__(self, num_correspondences, sigma): 11 | super().__init__() 12 | self.sigma = sigma 13 | self.num_correspondences = num_correspondences 14 | 15 | def forward(self, src_corr_points: Tensor, tgt_corr_points: Tensor) -> Tensor: 16 | sc_score_mat = spatial_consistency(src_corr_points, tgt_corr_points, self.sigma) # (N, N) 17 | corr_sc_counts = sc_score_mat.sum(dim=1) # (N) 18 | _, topk_indices = corr_sc_counts.topk(k=self.num_correspondences, largest=True) 19 | return topk_indices 20 | 21 | def extra_repr(self) -> str: 22 | param_strings = [f"num_correspondences={self.num_correspondences}", f"sigma={self.sigma:g}"] 23 | format_string = ", ".join(param_strings) 24 | return format_string 25 | -------------------------------------------------------------------------------- /vision3d/models/graph_pyramid/__init__.py: -------------------------------------------------------------------------------- 1 | from .graph_pyramid import build_grid_and_radius_graph_pyramid 2 | -------------------------------------------------------------------------------- /vision3d/models/graph_pyramid/graph_pyramid.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | import torch 4 | from torch import Tensor 5 | 6 | from vision3d.ops import grid_subsample_pack_mode, radius_search_pack_mode 7 | 8 | 9 | @torch.no_grad() 10 | def build_grid_and_radius_graph_pyramid( 11 | points: Tensor, 12 | lengths: Tensor, 13 | num_stages: int, 14 | voxel_size: float, 15 | search_radius: float, 16 | neighbor_limits: List[int], 17 | voxelize_first: bool = False, 18 | ) -> dict: 19 | points_list = [] 20 | lengths_list = [] 21 | neighbors_list = [] 22 | subsampling_list = [] 23 | upsampling_list = [] 24 | 25 | for i in range(num_stages): 26 | if voxelize_first or i > 0: 27 | points, lengths = grid_subsample_pack_mode(points, lengths, voxel_size) 28 | points_list.append(points) 29 | lengths_list.append(lengths) 30 | voxel_size *= 2.0 31 | 32 | for i in range(num_stages): 33 | cur_points = points_list[i] 34 | cur_lengths = lengths_list[i] 35 | 36 | neighbors = radius_search_pack_mode( 37 | cur_points, cur_points, cur_lengths, cur_lengths, search_radius, neighbor_limits[i] 38 | ) 39 | neighbors_list.append(neighbors) 40 | 41 | if i < num_stages - 1: 42 | sub_points = points_list[i + 1] 43 | sub_lengths = lengths_list[i + 1] 44 | 45 | subsampling = radius_search_pack_mode( 46 | sub_points, cur_points, sub_lengths, cur_lengths, search_radius, neighbor_limits[i] 47 | ) 48 | subsampling_list.append(subsampling) 49 | 50 | upsampling = radius_search_pack_mode( 51 | cur_points, sub_points, cur_lengths, sub_lengths, search_radius * 2, neighbor_limits[i + 1] 52 | ) 53 | upsampling_list.append(upsampling) 54 | 55 | search_radius *= 2 56 | 57 | return { 58 | "points": points_list, 59 | "lengths": lengths_list, 60 | "neighbors": neighbors_list, 61 | "subsampling": subsampling_list, 62 | "upsampling": upsampling_list, 63 | } 64 | -------------------------------------------------------------------------------- /vision3d/models/pat/__init__.py: -------------------------------------------------------------------------------- 1 | from .modules import * 2 | -------------------------------------------------------------------------------- /vision3d/models/point_transformer/__init__.py: -------------------------------------------------------------------------------- 1 | from vision3d.models.point_transformer.modules import PointTransformerBlock, TransitionDownBlock 2 | 3 | 4 | __all__ = [ 5 | "PointTransformerBlock", 6 | "TransitionDownBlock", 7 | ] 8 | -------------------------------------------------------------------------------- /vision3d/models/pointnet/__init__.py: -------------------------------------------------------------------------------- 1 | __all__ = [ 2 | 'TNet', 3 | 'TNetLoss', 4 | 'PointNetLoss', 5 | ] 6 | -------------------------------------------------------------------------------- /vision3d/models/pointnet2/__init__.py: -------------------------------------------------------------------------------- 1 | __all__ = [ 2 | "SetAbstractionModule", 3 | "MultiScaleSetAbstractionModule", 4 | "GlobalAbstractionModule", 5 | ] 6 | -------------------------------------------------------------------------------- /vision3d/models/predator/__init__.py: -------------------------------------------------------------------------------- 1 | from .modules import * 2 | 3 | __all__ = ['GCN'] 4 | -------------------------------------------------------------------------------- /vision3d/models/pri3d/__init__.py: -------------------------------------------------------------------------------- 1 | from .resunet import Res10UNet, Res18UNet, Res18UNetMultiRes, Res50UNet 2 | -------------------------------------------------------------------------------- /vision3d/models/siamese_transformer/__init__.py: -------------------------------------------------------------------------------- 1 | from .vanilla_siamese_transformer import VanillaSiameseTransformer 2 | -------------------------------------------------------------------------------- /vision3d/models/siamese_transformer/vanilla_siamese_transformer.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Tuple 2 | 3 | import torch.nn as nn 4 | from torch import Tensor 5 | 6 | from vision3d.layers import TransformerLayer 7 | 8 | 9 | class VanillaSiameseTransformer(nn.Module): 10 | def __init__( 11 | self, 12 | input_dim: int, 13 | output_dim: int, 14 | hidden_dim: int, 15 | num_heads: int, 16 | blocks: List[str], 17 | dropout: Optional[float] = None, 18 | activation_fn: str = "ReLU", 19 | ): 20 | super().__init__() 21 | 22 | self.in_proj = nn.Linear(input_dim, hidden_dim) 23 | self.out_proj = nn.Linear(hidden_dim, output_dim) 24 | 25 | self.blocks = blocks 26 | layers = [] 27 | for block in self.blocks: 28 | assert block in ["self", "cross"] 29 | layers.append(TransformerLayer(hidden_dim, num_heads, dropout=dropout, act_cfg=activation_fn)) 30 | self.transformer = nn.ModuleList(layers) 31 | 32 | def forward( 33 | self, 34 | src_feats: Tensor, 35 | tgt_feats: Tensor, 36 | src_masks: Optional[Tensor] = None, 37 | tgt_masks: Optional[Tensor] = None, 38 | ) -> Tuple[Tensor, Tensor]: 39 | src_tokens = self.in_proj(src_feats) 40 | tgt_tokens = self.in_proj(tgt_feats) 41 | 42 | for i, block in enumerate(self.blocks): 43 | if block == "self": 44 | src_tokens = self.transformer[i](src_tokens, src_tokens, src_tokens, k_masks=src_masks) 45 | tgt_tokens = self.transformer[i](tgt_tokens, tgt_tokens, tgt_tokens, k_masks=tgt_masks) 46 | else: 47 | src_tokens = self.transformer[i](src_tokens, tgt_tokens, tgt_tokens, k_masks=tgt_masks) 48 | tgt_tokens = self.transformer[i](tgt_tokens, src_tokens, src_tokens, k_masks=src_masks) 49 | 50 | src_feats = self.out_proj(src_tokens) 51 | tgt_feats = self.out_proj(tgt_tokens) 52 | 53 | return src_feats, tgt_feats 54 | -------------------------------------------------------------------------------- /vision3d/models/urr/__init__.py: -------------------------------------------------------------------------------- 1 | from .matching import CorrespondenceExtractor 2 | from .registration import RandomizedWeightedProcrustes 3 | from .render import DifferentiableRenderer 4 | -------------------------------------------------------------------------------- /vision3d/ops/back_project.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, Tuple, Union 2 | 3 | import torch 4 | from torch import Tensor 5 | 6 | 7 | def back_project( 8 | depth_mat: Tensor, 9 | intrinsics: Tensor, 10 | scaling_factor: float = 1000.0, 11 | depth_limit: Optional[float] = None, 12 | transposed: bool = False, 13 | return_mask: bool = False, 14 | ) -> Union[Tensor, Tuple[Tensor, Tensor]]: 15 | """Back project depth image to point cloud. 16 | 17 | Args: 18 | depth_mat (Tensor): the depth image in the shape of (B, H, W). 19 | intrinsics (Tensor): the intrinsic matrix in the shape of (B, 3, 3). 20 | scaling_factor (float): the depth scaling factor. Default: 1000. 21 | depth_limit (float, optional): ignore the pixels further than this value. 22 | transposed (bool): if True, the resulting point matrix is in the shape of (B, H, W, 3). 23 | return_mask (bool): if True, return a mask matrix where 0-depth points are False. Default: False. 24 | 25 | Returns: 26 | A Tensor of the point image in the shape of (B, 3, H, W). 27 | A Tensor of the mask image in the shape of (B, H, W). 28 | """ 29 | focal_x = intrinsics[..., 0:1, 0:1] 30 | focal_y = intrinsics[..., 1:2, 1:2] 31 | center_x = intrinsics[..., 0:1, 2:3] 32 | center_y = intrinsics[..., 1:2, 2:3] 33 | 34 | batch_size, height, width = depth_mat.shape 35 | coords = torch.arange(height * width).view(height, width).to(depth_mat.device).unsqueeze(0).expand_as(depth_mat) 36 | u = coords % width # (B, H, W) 37 | v = torch.div(coords, width, rounding_mode="floor") # (B, H, W) 38 | 39 | z = depth_mat / scaling_factor # (B, H, W) 40 | if depth_limit is not None: 41 | z.masked_fill_(torch.gt(z, depth_limit), 0.0) 42 | x = (u - center_x) * z / focal_x # (B, H, W) 43 | y = (v - center_y) * z / focal_y # (B, H, W) 44 | 45 | if transposed: 46 | points = torch.stack([x, y, z], dim=-1) # (B, H, W, 3) 47 | else: 48 | points = torch.stack([x, y, z], dim=1) # (B, 3, H, W) 49 | 50 | if not return_mask: 51 | return points 52 | 53 | masks = torch.gt(z, 0.0) 54 | 55 | return points, masks 56 | -------------------------------------------------------------------------------- /vision3d/ops/ball_query.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | 4 | from vision3d.utils.misc import load_ext 5 | 6 | 7 | ext_module = load_ext("vision3d.ext", ["ball_query"]) 8 | 9 | 10 | def ball_query( 11 | q_points: Tensor, 12 | s_points: Tensor, 13 | num_samples: int, 14 | radius: float, 15 | transposed: bool = False, 16 | ) -> Tensor: 17 | """Ball query. 18 | 19 | The first `num_samples` support points found in the local ball around each query points are returned. 20 | 21 | Args: 22 | q_points (Tensor): The query points in the shape of (B, N, 3). 23 | s_points (Tensor): The support points in the shape of (B, M, 3). 24 | num_samples (int): max neighbors. 25 | radius (float): search radius. 26 | transposed (bool=True): If True, the point tensor shapes are (B, 3, N) and (B, 3, M). 27 | 28 | Returns: 29 | A LongTensor of the indices of the sampled points in the shape of (B, M, K). 30 | """ 31 | if transposed: 32 | q_points = q_points.transpose(1, 2) # (B, 3, N) -> (B, N, 3) 33 | s_points = s_points.transpose(1, 2) # (B, 3, M) -> (B, M, 3) 34 | q_points = q_points.contiguous() 35 | s_points = s_points.contiguous() 36 | 37 | # (B, M, K) 38 | indices = torch.zeros(size=(q_points.shape[0], q_points.shape[1], num_samples), dtype=torch.long).cuda() 39 | 40 | ext_module.ball_query(q_points, s_points, indices, radius, num_samples) 41 | 42 | return indices 43 | -------------------------------------------------------------------------------- /vision3d/ops/contrib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qinzheng93/vision3d/7c61728eeeac76c9febf4f70c399358ce1a81b25/vision3d/ops/contrib/__init__.py -------------------------------------------------------------------------------- /vision3d/ops/cosine_similarity.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from torch.nn import functional as F 4 | 5 | 6 | def cosine_similarity(x: Tensor, y: Tensor, dim: int, normalized: bool = False) -> Tensor: 7 | """Cosine similarity between two set of feature vectors. 8 | 9 | The cosine similarity between two feature vector is defined as: 10 | 11 | s_{i} = \frac{\mathbf{x}_i \cdot \mathbf{y}_i}{\lVert \mathbf{x}_i \rVert \cdot \lVert \mathbf{y}_i \rVert}. 12 | 13 | The values of s_{i} are within [-1, 1]. We rescaled them to [0, 1] (larger is better): 14 | 15 | s*_{i} = 0.5 * (s_{i} + 1). 16 | 17 | Args: 18 | x (Tensor): (*, C, *) 19 | y (Tensor): (*, C, *) 20 | dim (int): the channel dim 21 | normalized (bool=False): If True, x and y is normalized. 22 | 23 | Returns: 24 | similarity (Tensor): (*) 25 | """ 26 | if not normalized: 27 | x = F.normalize(x, p=2, dim=dim) 28 | y = F.normalize(y, p=2, dim=dim) 29 | similarity = (x * y).sum(dim) 30 | similarity = 0.5 * (similarity + 1.0) 31 | return similarity 32 | 33 | 34 | def pairwise_cosine_similarity(x: Tensor, y: Tensor, normalized: bool = False, transposed: bool = False) -> Tensor: 35 | """Pairwise cosine similarity. 36 | 37 | The cosine similarity between two feature vector is defined as: 38 | 39 | s_{i, j} = \frac{\mathbf{x}_i \cdot \mathbf{y}_j}{\lVert \mathbf{x}_i \rVert \cdot \lVert \mathbf{y}_j \rVert}. 40 | 41 | The values of s_{i, j} are within [-1, 1]. We rescaled them to [0, 1] (larger is better): 42 | 43 | s*_{i, j} = 0.5 * (s_{i, j} + 1). 44 | 45 | Args: 46 | x (Tensor): (*, N, C) or (*, C, N) if transposed is True. 47 | y (Tensor): (*, M, C) or (*, C, M) if transposed is True. 48 | normalized (bool=False): If True, x and y is normalized. 49 | transposed (bool=False): If True, channel_dim is before length_dim. 50 | 51 | Returns: 52 | similarity_mat (Tensor): (*, N, M) 53 | """ 54 | if transposed: 55 | channel_dim = -2 56 | else: 57 | channel_dim = -1 58 | if not normalized: 59 | x = F.normalize(x, p=2, dim=channel_dim) 60 | y = F.normalize(y, p=2, dim=channel_dim) 61 | if transposed: 62 | similarity_mat = torch.matmul(x.transpose(-1, -2), y) 63 | else: 64 | similarity_mat = torch.matmul(x, y.transpose(-1, -2)) 65 | similarity_mat = 0.5 * (similarity_mat + 1.0) 66 | return similarity_mat 67 | -------------------------------------------------------------------------------- /vision3d/ops/dual_softmax.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import torch 4 | from torch import Tensor 5 | 6 | 7 | def dual_softmax(score_mat: Tensor, row_masks: Optional[Tensor] = None, col_masks: Optional[Tensor] = None) -> Tensor: 8 | """Dual softmax. 9 | 10 | output = softmax(input, dim=-1) * softmax(input, dim=-2) 11 | 12 | Args: 13 | score_mat (Tensor): the score matrix in the shape of (*, N, M). 14 | row_masks (BoolTensor, optional): the masks for the rows in the shape of (*, N). 15 | col_masks (BoolTensor, optional): the masks for the columns in the shape of (*, M). 16 | 17 | Returns: 18 | A Tensor of the dual softmax score matrix in the shape of (*, N, M). 19 | """ 20 | output_mat = torch.softmax(score_mat, dim=-1) * torch.softmax(score_mat, dim=-2) 21 | if row_masks is not None: 22 | output_mat = torch.where(row_masks.unsqueeze(-1), output_mat, torch.zeros_like(output_mat)) 23 | if col_masks is not None: 24 | output_mat = torch.where(col_masks.unsqueeze(-2), output_mat, torch.zeros_like(output_mat)) 25 | return output_mat 26 | -------------------------------------------------------------------------------- /vision3d/ops/eigenvector.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from torch import Tensor 4 | 5 | 6 | def leading_eigenvector(mat: Tensor, method: str = "power", num_iterations: int = 10) -> Tensor: 7 | """Compute the leading eigenvector of the matrix. 8 | 9 | Source: https://github.com/XuyangBai/PointDSC/blob/master/models/PointDSC.py 10 | 11 | Args: 12 | mat (Tensor): the input matrix in the shape of (*, M, M). 13 | method (str): the method to use: "power", "torch". Default: "power". 14 | num_iterations (int): the number of iterations for "power" method. Default: 10. 15 | 16 | Returns: 17 | A Tensor of the leading eigenvector. 18 | """ 19 | assert method in ["power", "torch"], f"Unsupported method: {method}." 20 | 21 | if method == "power": 22 | # power iteration algorithm 23 | leading_eig = torch.ones_like(mat[..., 0:1]) # (*, M, 1) 24 | leading_eig_last = leading_eig 25 | for i in range(num_iterations): 26 | leading_eig = torch.matmul(mat, leading_eig) 27 | leading_eig = F.normalize(leading_eig, p=2, dim=-2) 28 | if torch.allclose(leading_eig, leading_eig_last): 29 | break 30 | leading_eig_last = leading_eig 31 | leading_eig = leading_eig.squeeze(-1) 32 | else: 33 | _, eigenvectors = torch.symeig(mat, eigenvectors=True) 34 | leading_eig = eigenvectors[:, :, -1] 35 | 36 | return leading_eig 37 | -------------------------------------------------------------------------------- /vision3d/ops/furthest_point_sample.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | 4 | from vision3d.utils.misc import load_ext 5 | 6 | from .gather import gather 7 | 8 | 9 | ext_module = load_ext("vision3d.ext", ["furthest_point_sample"]) 10 | 11 | 12 | def furthest_point_sample( 13 | points: Tensor, num_samples: int, *input_tensors: Tensor, gather_points: bool = True, transposed: bool = False 14 | ): 15 | """ 16 | Furthest point sampling and gather the coordinates of the sampled points. 17 | 18 | Args: 19 | points (Tensor): The original points in shape of (B, N, 3). 20 | num_samples (int): The number of samples. 21 | input_tensors (List[Tensor]]): The features of the points (B, C, N). 22 | gather_points (bool=True): If True, gather points and features. 23 | transposed (bool=True): If True, the points shape is (B, 3, N). 24 | 25 | Returns: 26 | sampled_points (Tensor): The sampled points in shape of (B, M, 3) or (B, 3, M) if transposed. 27 | indices (LongTensor): The indices of the sample points in shape of (B, M). 28 | output_tensors (List[Tensor]): The list of sampled features of in shape of (B, C, M). 29 | """ 30 | if transposed: 31 | points = points.transpose(1, 2) # (B, 3, N) -> (B, N, 3) 32 | points = points.contiguous() 33 | 34 | distances = torch.full(size=(points.shape[0], points.shape[1]), fill_value=1e10).cuda() # (B, N) 35 | indices = torch.zeros(size=(points.shape[0], num_samples), dtype=torch.long).cuda() # (B, M) 36 | 37 | ext_module.furthest_point_sample(points, distances, indices, num_samples) 38 | 39 | if not gather_points: 40 | return indices 41 | 42 | points = points.transpose(1, 2).contiguous() # (B, N, 3) -> (B, 3, N) 43 | sampled_points = gather(points, indices) # (B, 3, M) 44 | if not transposed: 45 | sampled_points = sampled_points.transpose(1, 2).contiguous() # (B, 3, M) -> (B, M, 3) 46 | 47 | output_tensors = [gather(input_tensor, indices) for input_tensor in input_tensors] 48 | 49 | return sampled_points, *output_tensors 50 | -------------------------------------------------------------------------------- /vision3d/ops/gather.py: -------------------------------------------------------------------------------- 1 | from torch import Tensor 2 | from torch.autograd import Function 3 | 4 | from vision3d.utils.misc import load_ext 5 | 6 | 7 | ext_module = load_ext("vision3d.ext", ["gather_forward", "gather_backward"]) 8 | 9 | 10 | class GatherFunction(Function): 11 | @staticmethod 12 | def forward(ctx, inputs: Tensor, indices: Tensor, transposed: bool = False) -> Tensor: 13 | """Gather by indices. 14 | 15 | If not transposed: 16 | 17 | outputs[i][c][j] = inputs[i][c][indices[i][j]] 18 | 19 | If transposed: 20 | 21 | outputs[i][j][c] = inputs[i][indices[i][j]][c] 22 | 23 | Args: 24 | inputs (Tensor): The input tensor in the shape of (B, C, N). 25 | indices (Tensor): The indices to gather in the shape of (B, M). 26 | transposed (bool=False): If True, the inputs is in the shape of (B, N, C). 27 | 28 | Returns: 29 | A Tensor of the gathered inputs in the shape of (B, C, M) or (B, M, C) if transposed. 30 | """ 31 | if transposed: 32 | inputs = inputs.transpose(1, 2) # (B, N, C) -> (B, C, N) 33 | inputs = inputs.contiguous() 34 | indices = indices.contiguous() 35 | 36 | ctx.save_for_backward(indices) 37 | ctx.num_inputs = inputs.shape[-1] 38 | ctx.transposed = transposed 39 | 40 | outputs = inputs.new_zeros(size=(inputs.shape[0], inputs.shape[1], indices.shape[1])) # (B, C, M) 41 | 42 | ext_module.gather_forward(inputs, indices, outputs) 43 | 44 | if transposed: 45 | outputs = outputs.transpose(1, 2).contiguous() # (B, C, M) -> (B, M, C) 46 | 47 | return outputs 48 | 49 | @staticmethod 50 | def backward(ctx, grad_outputs: Tensor): 51 | indices = ctx.saved_tensors[0] 52 | num_inputs = ctx.num_inputs 53 | transposed = ctx.transposed 54 | 55 | if transposed: 56 | grad_outputs = grad_outputs.transpose(1, 2) # (B, M, C) -> (B, C, M) 57 | grad_outputs = grad_outputs.contiguous() 58 | 59 | # (B, C, N) 60 | grad_inputs = grad_outputs.new_zeros(size=(grad_outputs.shape[0], grad_outputs.shape[1], num_inputs)) 61 | 62 | ext_module.gather_backward(grad_outputs, indices, grad_inputs, num_inputs) 63 | 64 | if transposed: 65 | grad_inputs = grad_inputs.transpose(1, 2).contiguous() # (B, C, N) -> (B, N, C) 66 | 67 | return grad_inputs, None, None 68 | 69 | 70 | gather = GatherFunction.apply 71 | -------------------------------------------------------------------------------- /vision3d/ops/index_select.py: -------------------------------------------------------------------------------- 1 | from torch import Tensor 2 | 3 | 4 | def index_select(inputs: Tensor, indices: Tensor, dim: int) -> Tensor: 5 | """Advanced indices select. 6 | 7 | Returns a tensor `output` which indexes the `inputs` tensor along dimension `dim` using the entries in `indices` 8 | which is a `LongTensor`. 9 | 10 | Different from `torch.indices_select`, `indices` does not have to be 1-D. The `dim`-th dimension of `inputs` will 11 | be expanded to the number of dimensions in `indices`. 12 | 13 | For example, suppose the shape `inputs` is $(a_0, a_1, ..., a_{n-1})$, the shape of `indices` is 14 | $(b_0, b_1, ..., b_{m-1})$, and `dim` is $i$, then `output` is $(n+m-1)$-d tensor, whose shape is 15 | $(a_0, ..., a_{i-1}, b_0, b_1, ..., b_{m-1}, a_{i+1}, ..., a_{n-1})$. 16 | 17 | Args: 18 | inputs (Tensor): (a_0, a_1, ..., a_{n-1}) 19 | indices (LongTensor): (b_0, b_1, ..., b_{m-1}) 20 | dim (int): The dimension to index. 21 | 22 | Returns: 23 | output (Tensor): (a_0, ..., a_{dim-1}, b_0, ..., b_{m-1}, a_{dim+1}, ..., a_{n-1}) 24 | """ 25 | outputs = inputs.index_select(dim, indices.view(-1)) 26 | 27 | if indices.dim() > 1: 28 | if dim < 0: 29 | dim += inputs.dim() 30 | output_shape = inputs.shape[:dim] + indices.shape + inputs.shape[dim + 1 :] 31 | outputs = outputs.view(*output_shape) 32 | 33 | return outputs 34 | -------------------------------------------------------------------------------- /vision3d/ops/knn_points.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Union 2 | import torch 3 | from torch import Tensor 4 | 5 | from vision3d.utils.misc import load_ext 6 | 7 | 8 | ext_module = load_ext("vision3d.ext", ["knn_points"]) 9 | 10 | 11 | def knn_points( 12 | q_points: Tensor, s_points: Tensor, k: int, transposed: bool = False, return_distance: bool = False 13 | ) -> Union[Tensor, Tuple[Tensor, Tensor]]: 14 | """Heap sort based kNN search for point cloud. 15 | 16 | Args: 17 | q_points (Tensor): The query points in shape of (B, N, 3) or (B, 3, N) if transposed. 18 | s_points (Tensor): The support points in shape of (B, M, 3) or (B, 3, M) if transposed. 19 | k (int): The number of neighbors. 20 | transposed (bool=False): If True, the points are in shape of (B, 3, N). 21 | return_distance (bool=False): If True, return the distances of the kNN. 22 | 23 | Returns: 24 | A Tensor of the distances of the kNN in shape of (B, N, k). 25 | A LongTensor of the indices of the kNN in shape of (B, N, k). 26 | """ 27 | if transposed: 28 | q_points = q_points.transpose(1, 2) # (B, N, 3) -> (B, 3, N) 29 | s_points = s_points.transpose(1, 2) # (B, M, 3) -> (B, 3, M) 30 | q_points = q_points.contiguous() 31 | s_points = s_points.contiguous() 32 | 33 | knn_distances = q_points.new_zeros(size=(q_points.shape[0], q_points.shape[1], k)) # (B, N, k) 34 | knn_indices = torch.zeros(size=(q_points.shape[0], q_points.shape[1], k), dtype=torch.long).cuda() # (B, N, k) 35 | 36 | ext_module.knn(q_points, s_points, knn_distances, knn_indices, k) 37 | 38 | if return_distance: 39 | return knn_distances, knn_indices 40 | 41 | return knn_indices 42 | -------------------------------------------------------------------------------- /vision3d/ops/local_reference_frame.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn import functional as F 3 | 4 | from .group_gather import group_gather 5 | from .knn import knn 6 | 7 | 8 | def build_local_reference_frame(points, normals, num_neighbors, return_axes=False): 9 | """Build Local Reference Frame. 10 | 11 | Refer to ``The Perfect Match: 3D Point Cloud Matching with Smoothed Densities``. 12 | 13 | Args: 14 | points (Tensor): (B, 3, N) 15 | normals (Tensor): (B, 3, N) 16 | num_neighbors (int): number of neighbors. 17 | return_axes (bool=False): If True, return the lrf axes. 18 | 19 | Returns: 20 | x_axes (Tensor): (B, 3, N) 21 | y_axes (Tensor): (B, 3, N) 22 | z_axes (Tensor): (B, 3, N) 23 | knn_lrf_points (Tensor): (B, 3, N, K) 24 | """ 25 | knn_indices = knn(points, points, num_neighbors, transposed=True) # (B, N, k) 26 | knn_points = group_gather(points, knn_indices) # (B, 3, N, k) 27 | 28 | origins = points.unsqueeze(-1) # (B, 3, N, 1) 29 | z_axes = F.normalize(normals, dim=1, p=2).unsqueeze(-1) # (B, 3, N, 1) 30 | 31 | knn_offsets = knn_points - origins # (B, 3, N, k) 32 | knn_distances = torch.linalg.norm(knn_offsets, dim=1, keepdims=True) # (B, 1, N, k) 33 | knn_z_coords = (knn_offsets * z_axes).sum(dim=1, keepdims=True) # (B, 1, N, k) 34 | knn_plane_offsets = knn_offsets - knn_z_coords * z_axes # (B, 3, N, k) 35 | knn_weights = torch.abs(knn_z_coords) / (knn_distances + 1e-10) # (B, 1, N, k) 36 | knn_weights = knn_weights / knn_weights.sum(dim=-1, keepdims=True) # (B, 1, N, k) 37 | x_axes = F.normalize((knn_plane_offsets * knn_weights).sum(dim=-1, keepdims=True), dim=1, p=2) # (B, 3, N, 1) 38 | y_axes = torch.cross(x_axes, z_axes, dim=1) # (B, 3, N, 1) 39 | 40 | knn_x_coords = (knn_offsets * x_axes).sum(dim=1, keepdims=True) # (B, 1, N, k) 41 | knn_y_coords = (knn_offsets * y_axes).sum(dim=1, keepdims=True) # (B, 1, N, k) 42 | 43 | knn_lrf_points = torch.cat([knn_x_coords, knn_y_coords, knn_z_coords], dim=1) # (B, 3, N, k) 44 | 45 | if return_axes: 46 | x_axes = x_axes.squeeze(-1) 47 | y_axes = y_axes.squeeze(-1) 48 | z_axes = z_axes.squeeze(-1) 49 | return x_axes, y_axes, z_axes, knn_lrf_points 50 | else: 51 | return knn_lrf_points 52 | -------------------------------------------------------------------------------- /vision3d/ops/meshgrid.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def create_meshgrid(height: int, width: int, normalized: bool = False, flatten: bool = False, centering: bool = False): 5 | """Create meshgrid. 6 | 7 | Args: 8 | height (int): the height of the grid. 9 | width (int): the width of the grid. 10 | normalized (bool): if True, the resulting coordinates are normalized into [0, 1] inclusive. Default: False. 11 | flatten (bool): if True, the resulting coordinates are flattened. Default: False. 12 | centering (bool): if True, use the centers (+0.5) instead of the upperleft corners. Default: False. 13 | 14 | Returns: 15 | A tensor of the grid coordinates in the shape of (H, W, 2) if not flatten or (HxW, 2) otherwise. The tensor is 16 | float if normalized or int otherwise. 17 | """ 18 | if normalized and not centering: 19 | # TODO: use torch.arange in normalized non-centering implemention. 20 | # Currently, this affects 2D-3D registration models. 21 | h_values = torch.linspace(0.0, 1.0, steps=height).cuda() 22 | w_values = torch.linspace(0.0, 1.0, steps=width).cuda() 23 | else: 24 | h_values = torch.arange(height).cuda() 25 | w_values = torch.arange(width).cuda() 26 | if centering: 27 | h_values = h_values.float() + 0.5 28 | w_values = w_values.float() + 0.5 29 | if normalized: 30 | h_values = h_values.float() / float(height) 31 | w_values = w_values.float() / float(width) 32 | 33 | pixels = torch.cartesian_prod(h_values, w_values) # (H) x (W) -> (HxW, 2) 34 | 35 | if not flatten: 36 | pixels = pixels.view(height, width, 2) # (HxW, 2) -> (H, W, 2) 37 | return pixels 38 | -------------------------------------------------------------------------------- /vision3d/ops/nearest_interpolate.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .index_select import index_select 4 | 5 | 6 | def nearest_interpolate_pack_mode(inputs, neighbor_indices): 7 | """Pools features from the closest neighbors. 8 | 9 | WARNING: this function assumes the neighbors are ordered. 10 | 11 | Args: 12 | inputs: [n1, d] features matrix 13 | neighbor_indices: [n2, max_num] Only the first column is used for pooling 14 | 15 | Returns: 16 | outputs: [n2, d] pooled features matrix 17 | """ 18 | padded_inputs = torch.cat([inputs, torch.zeros_like(inputs[:1])], dim=0) 19 | outputs = index_select(padded_inputs, neighbor_indices[:, 0], dim=0) 20 | return outputs 21 | -------------------------------------------------------------------------------- /vision3d/ops/normal_estimation.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import torch 4 | from torch import Tensor 5 | 6 | from .index_select import index_select 7 | from .knn import knn 8 | 9 | 10 | def estimate_normals( 11 | q_points: Tensor, s_points: Optional[Tensor] = None, k: int = 50, disambiguate_directions: bool = True 12 | ): 13 | """Estimate normals for the query points from the support points. 14 | 15 | Modified from [PyTorch3d](https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/ops/points_normals.py). 16 | 17 | Args: 18 | q_points (Tensor): query points (N, 3) 19 | s_points (Tensor=None): support points. If None, use q_points (M, 3) 20 | k (int=50): number of neighbors 21 | disambiguate_directions (bool=True): disambiguate directions. 22 | 23 | Returns: 24 | q_normals (Tensor): normals of the query points (N, 3) 25 | """ 26 | if s_points is None: 27 | s_points = q_points 28 | 29 | knn_indices = knn(q_points, s_points, k) # (N, K) 30 | knn_points = index_select(s_points, knn_indices, dim=0) # (N, K, 3) 31 | centroids = knn_points.mean(dim=1, keepdim=True) # (N, 1, 3) 32 | knn_offsets = knn_points - centroids # (N, K, 3) 33 | 34 | cov_mat = knn_offsets.unsqueeze(3) * knn_offsets.unsqueeze(2) # (N, K, 3, 1) x (N, K, 1, 3) 35 | cov_mat = cov_mat.mean(dim=1) # (N, 3, 3) 36 | _, eigenvectors = torch.symeig(cov_mat.cpu(), eigenvectors=True) # (N, 3), (N, 3, 3) 37 | eigenvectors = eigenvectors.cuda() 38 | normals = eigenvectors[:, :, 0] # (N, 3) 39 | 40 | if disambiguate_directions: 41 | knn_offsets = knn_points - q_points.unsqueeze(1) # (N, K, 3) - (N, 3) -> (N, 1, 3) 42 | projections = (normals.unsqueeze(1) * knn_offsets).sum(2) # (N, 3) -> (N, 1, 3) x (N, K, 3) -> (N, K) 43 | flips = torch.gt(projections.sum(dim=-1, keepdim=True), 0).float() # (N, K) -> (N, 1) 44 | # flips = torch.lt(torch.gt(projections, 0).float().mean(dim=1, keepdim=True), 0.5).float() # (N, K) -> (N, 1) 45 | normals = (1.0 - 2.0 * flips) * normals 46 | 47 | return normals 48 | -------------------------------------------------------------------------------- /vision3d/ops/numeric.py: -------------------------------------------------------------------------------- 1 | from typing import Union 2 | 3 | import torch 4 | from torch import Tensor 5 | 6 | 7 | def safe_divide(a: Union[Tensor, float], b: Tensor, eps: float = 1e-6): 8 | b = torch.clamp(b, min=eps) 9 | return a / b 10 | 11 | 12 | def safe_sqrt(a: Tensor, eps: float = 1e-6): 13 | a = torch.clamp(a, min=eps) 14 | return torch.sqrt(a) 15 | -------------------------------------------------------------------------------- /vision3d/ops/pairwise_distance.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | 4 | 5 | def pairwise_distance( 6 | x: Tensor, 7 | y: Tensor, 8 | normalized: bool = False, 9 | transposed: bool = False, 10 | squared: bool = True, 11 | strict: bool = False, 12 | eps: float = 1e-8, 13 | ) -> Tensor: 14 | """Pairwise distance of two (batched) point clouds. 15 | 16 | Args: 17 | x (Tensor): the row tensor in the shape of (*, N, C). 18 | y (Tensor): the column tensor in the shape of (*, M, C). 19 | normalized (bool): if the points are normalized, we have "x^2 + y^2 = 1", so "d^2 = 2 - 2xy". Default: False. 20 | transposed (bool): if True, x and y are in the shapes of (*, C, N) and (*, C, M) respectively. Default: False. 21 | squared (bool): if True, return squared distance. Default: True. 22 | strict (bool): if True, use strict mode to guarantee precision. Default: False. 23 | eps (float): a safe number for sqrt. Default: 1e-8. 24 | 25 | Returns: 26 | dist: Tensor (*, N, M) 27 | """ 28 | if strict: 29 | if transposed: 30 | displacements = x.unsqueeze(-1) - y.unsqueeze(-2) # (*, C, N, 1) x (*, C, 1, M) -> (*, C, N, M) 31 | distances = torch.linalg.norm(displacements, dim=-3) # (*, C, N, M) -> (*, N, M) 32 | else: 33 | displacements = x.unsqueeze(-2) - y.unsqueeze(-3) # (*, N, 1, C) x (*, 1, M, C) -> (*, N, M, C) 34 | distances = torch.linalg.norm(displacements, dim=-1) # (*, N, M, C) -> (*, N, M) 35 | 36 | if squared: 37 | distances = distances.pow(2) 38 | else: 39 | if transposed: 40 | channel_dim = -2 41 | xy = torch.matmul(x.transpose(-1, -2), y) # [(*, C, N) -> (*, N, C)] x (*, C, M) 42 | else: 43 | channel_dim = -1 44 | xy = torch.matmul(x, y.transpose(-1, -2)) # (*, N, C) x [(*, M, C) -> (*, C, M)] 45 | 46 | if normalized: 47 | distances = 2.0 - 2.0 * xy 48 | else: 49 | x2 = torch.sum(x ** 2, dim=channel_dim).unsqueeze(-1) # (*, N, C) or (*, C, N) -> (*, N) -> (*, N, 1) 50 | y2 = torch.sum(y ** 2, dim=channel_dim).unsqueeze(-2) # (*, M, C) or (*, C, M) -> (*, M) -> (*, 1, M) 51 | distances = x2 - 2 * xy + y2 52 | 53 | distances = distances.clamp(min=0.0) 54 | 55 | if not squared: 56 | distances = torch.sqrt(distances + eps) 57 | 58 | return distances 59 | -------------------------------------------------------------------------------- /vision3d/ops/pooling.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .index_select import index_select 4 | 5 | 6 | def local_maxpool_pack_mode(feats, neighbor_indices): 7 | """Max pooling from neighbors in pack mode. 8 | 9 | Args: 10 | feats (Tensor): The input features in the shape of (N, C). 11 | neighbor_indices (LongTensor): The neighbor indices in the shape of (M, K). 12 | 13 | Returns: 14 | pooled_feats (Tensor): The pooled features in the shape of (M, C). 15 | """ 16 | feats = torch.cat((feats, torch.zeros_like(feats[:1, :])), 0) # (N+1, C) 17 | neighbor_feats = index_select(feats, neighbor_indices, dim=0) # (M, K, C) 18 | pooled_feats = neighbor_feats.max(1)[0] # (M, K) 19 | return pooled_feats 20 | 21 | 22 | def global_avgpool_pack_mode(feats, lengths): 23 | """Global average pooling over batch. 24 | 25 | Args: 26 | feats (Tensor): The input features in the shape of (N, C). 27 | lengths (LongTensor): The length of each sample in the batch in the shape of (B). 28 | 29 | Returns: 30 | feats (Tensor): The pooled features in the shape of (B, C). 31 | """ 32 | feats_list = [] 33 | start_index = 0 34 | for batch_index in range(lengths.shape[0]): 35 | end_index = start_index + lengths[batch_index].item() 36 | feats_list.append(torch.mean(feats[start_index:end_index], dim=0)) 37 | start_index = end_index 38 | feats = torch.stack(feats_list, dim=0) 39 | return feats 40 | -------------------------------------------------------------------------------- /vision3d/ops/random_point_sample.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | 4 | from .gather import gather 5 | 6 | 7 | def random_point_sample(points: Tensor, num_samples: int, gather_points: bool = True): 8 | batch_size, _, num_points = points.shape 9 | weights = torch.ones(size=(batch_size, num_points)).cuda() 10 | indices = torch.multinomial(weights, num_samples) 11 | if gather_points: 12 | samples = gather(points, indices) 13 | return samples 14 | else: 15 | return indices 16 | -------------------------------------------------------------------------------- /vision3d/ops/random_sample.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, Union 2 | 3 | import numpy as np 4 | import torch 5 | from torch import Tensor 6 | 7 | 8 | def random_sample_from_scores(scores: Tensor, size: int, replace: bool = False) -> Tensor: 9 | """Random sample with `scores` as probability. 10 | 11 | Args: 12 | scores: Tensor (N, *) 13 | size: int 14 | replace: bool 15 | 16 | Returns: 17 | sel_indices: torch.LongTensor (size,) 18 | """ 19 | if scores.shape[0] > size or replace: 20 | probs = scores / scores.sum() 21 | probs = probs.detach().cpu().numpy() 22 | sel_indices = np.random.choice(scores.shape[0], size=size, p=probs, replace=replace) 23 | sel_indices = torch.from_numpy(sel_indices).cuda() 24 | else: 25 | sel_indices = torch.arange(scores.shape[0]).cuda() 26 | return sel_indices 27 | 28 | 29 | def random_choice( 30 | a: Union[Tensor, int], 31 | size: int = None, 32 | replace: bool = True, 33 | p: Optional[Tensor] = None, 34 | ) -> Tensor: 35 | """Numpy-style random choice. 36 | 37 | Args: 38 | a (Tensor|int): 1-D tensor or int. 39 | size (int=None): the number of choices. 40 | replace (bool=True): True if sample with replacement. 41 | p (Tensor=None): probabilities. 42 | 43 | Returns: 44 | selected (Tensor): selected items if ``a`` is a tensor or selected indices if ``a`` is int. 45 | """ 46 | assert isinstance(a, (Tensor, int)), "'a' must be a 1-D tensor or int." 47 | if p is not None: 48 | assert isinstance(p, Tensor) and p.dim() == 1, "'p' must be a 1-D tensor or None." 49 | p = p.detach().cpu().numpy() 50 | if isinstance(a, Tensor): 51 | assert a.dim() == 1, "'a' must be a 1-D tensor or int." 52 | sel_indices = np.random.choice(a.shape[0], size=size, replace=replace, p=p) 53 | sel_indices = torch.from_numpy(sel_indices).cuda() 54 | return a[sel_indices] 55 | else: 56 | if p is not None: 57 | assert a == p.shape[0], f"'a' must be equal to the length of 'p' ({a} vs {p.shape[0]})." 58 | sel_indices = np.random.choice(a, size=size, replace=replace, p=p) 59 | sel_indices = torch.from_numpy(sel_indices).cuda() 60 | return sel_indices 61 | -------------------------------------------------------------------------------- /vision3d/ops/three_nn.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple 2 | 3 | import torch 4 | from torch import Tensor 5 | 6 | from vision3d.utils.misc import load_ext 7 | 8 | 9 | ext_module = load_ext("vision3d.ext", ["three_nn"]) 10 | 11 | 12 | def three_nn(q_points: Tensor, s_points: Tensor, transposed: bool = False) -> Tuple[Tensor, Tensor]: 13 | """ 14 | Compute the three nearest neighbors for the query points in the support points. 15 | 16 | Args: 17 | q_points (Tensor): The query points in shape of (B, N, 3) or (B, 3, N) if transposed. 18 | s_points (Tensor): The support points in shape of (B, M, 3) or (B, 3, M) if transposed. 19 | transposed (bool=True): If True, the points shape is (B, 3, N). 20 | 21 | Returns: 22 | A Tensor of the distances of the 3-NN in the shape of (B, N, 3). 23 | A LongTensor of the indices of the 3-NN in the shape of (B, N, 3). 24 | """ 25 | if transposed: 26 | q_points = q_points.transpose(1, 2) # (B, N, 3) -> (B, 3, N) 27 | s_points = s_points.transpose(1, 2) # (B, M, 3) -> (B, 3, M) 28 | q_points = q_points.contiguous() 29 | s_points = s_points.contiguous() 30 | 31 | assert s_points.shape[1] >= 3, "s_points must have at least 3 points." 32 | 33 | tnn_distances = q_points.new_zeros(size=(q_points.shape[0], q_points.shape[1], 3)) # (B, N, 3) 34 | tnn_indices = torch.zeros(size=(q_points.shape[0], q_points.shape[1], 3), dtype=torch.long).cuda() # (B, N, 3) 35 | 36 | ext_module.three_nn(q_points, s_points, tnn_distances, tnn_indices) 37 | 38 | return tnn_distances, tnn_indices 39 | -------------------------------------------------------------------------------- /vision3d/ops/vector_angle.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | import numpy as np 4 | 5 | 6 | def rad2deg(rad: Tensor) -> Tensor: 7 | factor = 180.0 / np.pi 8 | deg = rad * factor 9 | return deg 10 | 11 | 12 | def deg2rad(deg: Tensor) -> Tensor: 13 | factor = np.pi / 180.0 14 | rad = deg * factor 15 | return rad 16 | 17 | 18 | def vector_angle(x: Tensor, y: Tensor, dim: int, use_degree: bool = False) -> Tensor: 19 | """Compute the angles between two set of 3D vectors. 20 | 21 | Args: 22 | x (Tensor): set of vectors (*, 3, *) 23 | y (Tensor): set of vectors (*, 3, *). 24 | dim (int): dimension index of the coordinates. 25 | use_degree (bool=False): If True, return angles in degree instead of rad. 26 | 27 | Returns: 28 | angle (Tensor): (*) 29 | """ 30 | cross = torch.linalg.norm(torch.cross(x, y, dim=dim), dim=dim) # (*, 3 *) x (*, 3, *) -> (*, 3, *) -> (*) 31 | dot = torch.sum(x * y, dim=dim) # (*, 3 *) x (*, 3, *) -> (*) 32 | angle = torch.atan2(cross, dot) # (*) 33 | if use_degree: 34 | angle = rad2deg(angle) 35 | return angle 36 | -------------------------------------------------------------------------------- /vision3d/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qinzheng93/vision3d/7c61728eeeac76c9febf4f70c399358ce1a81b25/vision3d/utils/__init__.py -------------------------------------------------------------------------------- /vision3d/utils/average_meter.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | class AverageMeter: 5 | def __init__(self): 6 | self._records = [] 7 | 8 | def update(self, result): 9 | self._records.append(result) 10 | 11 | def reset(self): 12 | self._records.clear() 13 | 14 | def get_records(self, last_n=None): 15 | if last_n is not None: 16 | return self._records[-last_n:] 17 | return self._records 18 | 19 | def sum(self, last_n=None): 20 | records = self.get_records(last_n=last_n) 21 | return np.sum(records) 22 | 23 | def mean(self, last_n=None): 24 | records = self.get_records(last_n=last_n) 25 | return np.mean(records) 26 | 27 | def std(self, last_n=None): 28 | records = self.get_records(last_n=last_n) 29 | return np.std(records) 30 | 31 | def median(self, last_n=None): 32 | records = self.get_records(last_n=last_n) 33 | return np.median(records) 34 | -------------------------------------------------------------------------------- /vision3d/utils/distributed.py: -------------------------------------------------------------------------------- 1 | import os 2 | from functools import wraps 3 | 4 | import torch 5 | import torch.distributed as dist 6 | from torch import Tensor 7 | 8 | # distributed env 9 | 10 | 11 | def setup_distributed(): 12 | if "LOCAL_RANK" in os.environ: 13 | local_rank = int(os.environ["LOCAL_RANK"]) 14 | torch.cuda.set_device(local_rank) 15 | dist.init_process_group(backend="nccl") 16 | else: 17 | # single-gpu mode, use cuda:0 18 | torch.cuda.set_device(0) 19 | 20 | 21 | def is_distributed() -> bool: 22 | if not dist.is_available(): 23 | return False 24 | if not dist.is_initialized(): 25 | return False 26 | return True 27 | 28 | 29 | def get_world_size() -> int: 30 | if not is_distributed(): 31 | return 1 32 | return dist.get_world_size() 33 | 34 | 35 | def get_local_rank() -> int: 36 | if not is_distributed(): 37 | return 0 38 | return dist.get_rank() 39 | 40 | 41 | def is_master() -> bool: 42 | return get_local_rank() == 0 43 | 44 | 45 | def master_only(func): 46 | @wraps(func) 47 | def wrapper(*args, **kwargs): 48 | if is_master(): 49 | return func(*args, **kwargs) 50 | 51 | return wrapper 52 | 53 | 54 | # reduce tensor 55 | 56 | 57 | def all_reduce_tensor(tensor, world_size=None): 58 | """Average reduce a tensor across all workers.""" 59 | if world_size is None: 60 | world_size = get_world_size() 61 | reduced_tensor = tensor.clone() 62 | dist.all_reduce(reduced_tensor) 63 | reduced_tensor /= world_size 64 | return reduced_tensor 65 | 66 | 67 | def all_reduce_tensors(x, world_size=None): 68 | """Average reduce all tensors across all workers.""" 69 | if isinstance(x, list): 70 | x = [all_reduce_tensors(item, world_size=world_size) for item in x] 71 | elif isinstance(x, tuple): 72 | x = tuple([all_reduce_tensors(item, world_size=world_size) for item in x]) 73 | elif isinstance(x, dict): 74 | x = {key: all_reduce_tensors(value, world_size=world_size) for key, value in x.items()} 75 | elif isinstance(x, Tensor): 76 | x = all_reduce_tensor(x, world_size=world_size) 77 | return x 78 | -------------------------------------------------------------------------------- /vision3d/utils/misc.py: -------------------------------------------------------------------------------- 1 | import importlib 2 | from functools import wraps 3 | 4 | import torch 5 | 6 | # deprecate decorator 7 | 8 | 9 | def deprecated(replacement: str): 10 | def decorate(func): 11 | @wraps(func) 12 | def wrapper(*args, **kwargs): 13 | raise RuntimeError(f"{func.__name__} is deprecated. Use {replacement} instead.") 14 | 15 | return wrapper 16 | 17 | return decorate 18 | 19 | 20 | # importlib utilities 21 | 22 | 23 | def load_ext(name, functions): 24 | ext_module = importlib.import_module(name) 25 | for function in functions: 26 | assert hasattr(ext_module, function), f"Function '{function}' missing in '{name}'." 27 | return ext_module 28 | 29 | 30 | # log utilities 31 | 32 | 33 | def get_print_format(value): 34 | if isinstance(value, (int, str)): 35 | return "" 36 | if value == 0: 37 | return ".3f" 38 | if value < 1e-5: 39 | return ".3e" 40 | if value < 1e-2: 41 | return ".6f" 42 | return ".3f" 43 | 44 | 45 | def get_format_strings(result_dict): 46 | """Get format string for a list of key-value pairs.""" 47 | format_strings = [] 48 | if "metadata" in result_dict: 49 | # handle special key "metadata" 50 | format_strings.append(result_dict["metadata"]) 51 | for key, value in result_dict.items(): 52 | if key == "metadata": 53 | continue 54 | format_string = f"{key}: {value:{get_print_format(value)}}" 55 | format_strings.append(format_string) 56 | return format_strings 57 | 58 | 59 | def get_log_string( 60 | result_dict, epoch=None, max_epoch=None, iteration=None, max_iteration=None, lr=None, time_dict=None 61 | ): 62 | log_strings = [] 63 | if epoch is not None: 64 | epoch_string = f"epoch: {epoch}" 65 | if max_epoch is not None: 66 | epoch_string += f"/{max_epoch}" 67 | log_strings.append(epoch_string) 68 | if iteration is not None: 69 | iter_string = f"iter: {iteration}" 70 | if max_iteration is not None: 71 | iter_string += f"/{max_iteration}" 72 | log_strings.append(iter_string) 73 | log_strings += get_format_strings(result_dict) 74 | if lr is not None: 75 | log_strings.append("lr: {:.3e}".format(lr)) 76 | if time_dict is not None: 77 | time_string = "time: " + "/".join([f"{time_dict[key]:.3f}s" for key in time_dict]) 78 | log_strings.append(time_string) 79 | message = ", ".join(log_strings) 80 | return message 81 | -------------------------------------------------------------------------------- /vision3d/utils/opencv.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import cv2 4 | import numpy as np 5 | from numpy import ndarray 6 | 7 | from vision3d.array_ops import axis_angle_to_rotation_matrix, get_transform_from_rotation_translation 8 | 9 | 10 | def registration_with_pnp_ransac( 11 | corr_points: ndarray, 12 | corr_pixels: ndarray, 13 | intrinsics: ndarray, 14 | distortion: Optional[ndarray] = None, 15 | num_iterations: int = 5000, 16 | distance_tolerance: float = 8.0, 17 | transposed: bool = True, 18 | ) -> ndarray: 19 | """PnP-RANSAC registration with OpenCV. 20 | 21 | Note: 22 | 1. cv2.solvePnPRansac() requires the pixels are in the order of (w, h). 23 | 24 | Args: 25 | corr_points (array): a float array of the 3D correspondence points in the shape of (N, 3). 26 | corr_pixels (array): an int array of the 2D correspondence pixels in the shape of (N, 2). 27 | intrinsics (array): a float array of the camera intrinsics in the shape of (3, 3). 28 | distortion (array, optional): a float array of the distortion parameter in the shape of (4, 1) or (12, 1). 29 | num_iterations (int): the number of ransac iterations. 30 | distance_tolerance (float): the distance tolerance for ransac. 31 | transposed (bool): if True, the pixel coordinates are in the order of (h, w) or (w, h) otherwise. 32 | 33 | Returns: 34 | A float array of the estimated transformation from 3D to 2D. 35 | """ 36 | if corr_points.shape[0] < 4: 37 | # too few correspondences, return None 38 | return None 39 | 40 | if distortion is None: 41 | distortion = np.zeros(shape=(4, 1)) 42 | if transposed: 43 | corr_pixels = np.stack([corr_pixels[..., 1], corr_pixels[..., 0]], axis=-1) # (h, w) -> (w, h) 44 | 45 | corr_points = corr_points.astype(np.float) 46 | corr_pixels = corr_pixels.astype(np.float) 47 | intrinsics = intrinsics.astype(np.float) 48 | 49 | _, axis_angle, translation, _ = cv2.solvePnPRansac( 50 | corr_points, 51 | corr_pixels, 52 | intrinsics, 53 | distortion, 54 | iterationsCount=num_iterations, 55 | reprojectionError=distance_tolerance, 56 | flags=cv2.SOLVEPNP_P3P, 57 | ) 58 | axis_angle = axis_angle[:, 0] 59 | translation = translation[:, 0] 60 | rotation = axis_angle_to_rotation_matrix(axis_angle) 61 | estimated_transform = get_transform_from_rotation_translation(rotation, translation) 62 | 63 | return estimated_transform 64 | -------------------------------------------------------------------------------- /vision3d/utils/parser.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | _PARSER = None 4 | 5 | 6 | def get_default_parser(): 7 | global _PARSER 8 | if _PARSER is None: 9 | _PARSER = argparse.ArgumentParser() 10 | return _PARSER 11 | 12 | 13 | def parse_args(): 14 | parser = get_default_parser() 15 | args = parser.parse_args() 16 | return args 17 | 18 | 19 | def add_trainer_args(): 20 | parser = get_default_parser() 21 | parser.add_argument_group("trainer", "trainer arguments") 22 | parser.add_argument("--resume", action="store_true", help="resume training from the latest checkpoint") 23 | parser.add_argument("--log_steps", type=int, default=10, help="logging steps") 24 | parser.add_argument("--debug", action="store_true", help="debug mode with grad check") 25 | parser.add_argument("--detect_anomaly", action="store_true", help="detect anomaly with autograd") 26 | parser.add_argument("--save_latest_n_models", type=int, default=-1, help="save latest n models") 27 | parser.add_argument("--cudnn_deterministic", type=bool, default=True, help="use deterministic method") 28 | 29 | 30 | def add_tester_args(): 31 | parser = get_default_parser() 32 | parser.add_argument_group("tester", "tester arguments") 33 | parser.add_argument("--checkpoint", default=None, help="load from checkpoint") 34 | parser.add_argument("--test_epoch", type=int, default=None, help="test epoch") 35 | parser.add_argument("--cudnn_deterministic", type=bool, default=True, help="use deterministic method") 36 | -------------------------------------------------------------------------------- /vision3d/utils/profiling.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | import torch 4 | 5 | 6 | class CudaRuntimeProfiler: 7 | def __init__(self, description): 8 | self.description = description 9 | self.last_time = 0.0 10 | 11 | def __enter__(self): 12 | torch.cuda.synchronize() 13 | self.last_time = time.time() 14 | return self 15 | 16 | def __exit__(self, exc_type, exc_value, traceback): 17 | torch.cuda.synchronize() 18 | duration = time.time() - self.last_time 19 | print(f"{self.description}: {duration:.3f}s") 20 | 21 | 22 | class CpuRuntimeProfiler: 23 | def __init__(self, description): 24 | self.description = description 25 | self.last_time = 0.0 26 | 27 | def __enter__(self): 28 | self.last_time = time.time() 29 | return self 30 | 31 | def __exit__(self, exc_type, exc_value, traceback): 32 | duration = time.time() - self.last_time 33 | print(f"{self.description}: {duration:.3f}s") 34 | 35 | 36 | def profile_cpu_runtime(description): 37 | """Profile CPU runtime. 38 | 39 | Usage: 40 | 41 | ``` 42 | with profile_cpu_runtime("call function"): 43 | outputs = function_call(inputs) 44 | ``` 45 | 46 | Args: 47 | description (str): The description string for the code. 48 | 49 | Returns: 50 | A CpuRuntimeProfiler object. 51 | """ 52 | return CpuRuntimeProfiler(description) 53 | 54 | 55 | def profile_cuda_runtime(description): 56 | """Profile Cuda runtime. 57 | 58 | Usage: 59 | 60 | ``` 61 | with profile_cuda_runtime("call function"): 62 | outputs = function_call(inputs) 63 | ``` 64 | 65 | The CUDA device is synchronized during profiling. 66 | 67 | Args: 68 | description (str): The description string for the code. 69 | 70 | Returns: 71 | A CudaRuntimeProfiler object. 72 | """ 73 | return CudaRuntimeProfiler(description) 74 | -------------------------------------------------------------------------------- /vision3d/utils/tensor.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from torch import Tensor 4 | 5 | 6 | def tensor_to_array(x): 7 | """Release all pytorch tensors to item or numpy arrays.""" 8 | if isinstance(x, list): 9 | x = [tensor_to_array(item) for item in x] 10 | elif isinstance(x, tuple): 11 | x = tuple([tensor_to_array(item) for item in x]) 12 | elif isinstance(x, dict): 13 | x = {key: tensor_to_array(value) for key, value in x.items()} 14 | elif isinstance(x, Tensor): 15 | if x.numel() == 1: 16 | x = x.item() 17 | else: 18 | x = x.detach().cpu().numpy() 19 | return x 20 | 21 | 22 | def array_to_tensor(x): 23 | """Convert all numpy arrays to pytorch tensors.""" 24 | if isinstance(x, list): 25 | x = [array_to_tensor(item) for item in x] 26 | elif isinstance(x, tuple): 27 | x = tuple([array_to_tensor(item) for item in x]) 28 | elif isinstance(x, dict): 29 | x = {key: array_to_tensor(value) for key, value in x.items()} 30 | elif isinstance(x, np.ndarray): 31 | x = torch.from_numpy(x) 32 | return x 33 | 34 | 35 | def move_to_cpu(x): 36 | """Move all tensors to cpu.""" 37 | if isinstance(x, list): 38 | x = [move_to_cpu(item) for item in x] 39 | elif isinstance(x, tuple): 40 | x = tuple([move_to_cpu(item) for item in x]) 41 | elif isinstance(x, dict): 42 | x = {key: move_to_cpu(value) for key, value in x.items()} 43 | elif isinstance(x, Tensor): 44 | x = x.cpu() 45 | return x 46 | 47 | 48 | def move_to_cuda(x): 49 | """Move all tensors to cuda.""" 50 | if isinstance(x, list): 51 | x = [move_to_cuda(item) for item in x] 52 | elif isinstance(x, tuple): 53 | x = tuple([move_to_cuda(item) for item in x]) 54 | elif isinstance(x, dict): 55 | x = {key: move_to_cuda(value) for key, value in x.items()} 56 | elif isinstance(x, Tensor): 57 | x = x.cuda() 58 | return x 59 | -------------------------------------------------------------------------------- /vision3d/utils/timer.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | class Timer: 5 | def __init__(self): 6 | self._total_time = {} 7 | self._count_time = {} 8 | self._last_time = {} 9 | self._keys = [] 10 | 11 | def register_timer(self, key): 12 | self._total_time[key] = 0.0 13 | self._count_time[key] = 0 14 | self._last_time[key] = None 15 | self._keys.append(key) 16 | 17 | def tic(self, key): 18 | if key not in self._keys: 19 | self.register_timer(key) 20 | self._last_time[key] = time.time() 21 | 22 | def toc(self, key): 23 | assert key in self._keys 24 | assert self._last_time[key] is not None, "'tic' must be called before 'toc'." 25 | duration = time.time() - self._last_time[key] 26 | self._total_time[key] += duration 27 | self._count_time[key] += 1 28 | self._last_time[key] = None 29 | 30 | def get_time(self, key): 31 | assert key in self._keys 32 | return self._total_time[key] / (self._count_time[key] + 1e-12) 33 | 34 | def tostring(self, keys=None, verbose=True): 35 | if keys is None: 36 | keys = self._keys 37 | if verbose: 38 | log_strings = [f"{key}: {self.get_time(key):.3f}s" for key in keys if key in self._keys] 39 | format_string = ", ".join(log_strings) 40 | else: 41 | log_strings = [f"{self.get_time(key):.3f}s" for key in keys if key in self._keys] 42 | format_string = "time: " + "/".join(log_strings) 43 | return format_string 44 | 45 | def summary(self, keys=None): 46 | if keys is None: 47 | keys = self._keys 48 | summary_dict = {key: self.get_time(key) for key in keys} 49 | return summary_dict 50 | --------------------------------------------------------------------------------