Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-09-15 14:13:49 +08:00
commit 7e4ef6f98b
7002 changed files with 1452249 additions and 0 deletions
@@ -0,0 +1,137 @@
+/***********************************************************************
+ * The code in this file is modified from the original code in the
+ * https://github.com/jlblancoc/nanoflann repository.
+ *
+ * Software License Agreement (BSD License)
+ *
+ * Copyright 2008-2009  Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
+ * Copyright 2008-2009  David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
+ * Copyright 2011-2025  Jose Luis Blanco (joseluisblancoc@gmail.com).
+ *   All rights reserved.
+ *
+ * THE BSD LICENSE
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+ * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+ * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
+ * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
+ * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *************************************************************************/
+
+#ifndef FLANN_CUDA_KD_TREE_BUILDER_H_
+#define FLANN_CUDA_KD_TREE_BUILDER_H_
+
+#include <thrust/host_vector.h>
+#include <thrust/device_vector.h>
+#include <thrust/sort.h>
+#include <thrust/partition.h>
+#include <thrust/unique.h>
+#include <thrust/scan.h>
+#include "math_util.h"
+#include <cstdlib>
+#include <map>
+
+namespace tinyflann {
+
+// Distance types.
+struct CudaL1;
+struct CudaL2;
+
+// Parameters
+struct KDTreeCuda3dIndexParams {
+    int leaf_max_size = 64;
+};
+
+struct SearchParams {
+    SearchParams(int checks_ = 32, float eps_ = 0.0, bool sorted_ = true)
+        : checks(checks_), eps(eps_), sorted(sorted_) {
+        max_neighbors = -1;
+        use_heap = true;
+    }
+
+    // how many leafs to visit when searching for neighbours (-1 for unlimited)
+    int checks;
+    // search for eps-approximate neighbours (default: 0)
+    float eps;
+    // only for radius search, require neighbours sorted by distance (default: true)
+    bool sorted;
+    // maximum number of neighbors radius search should return (-1 for unlimited)
+    int max_neighbors;
+    // use a heap to manage the result set (default: FLANN_Undefined)
+    bool use_heap;
+};
+
+template <typename Distance>
+class KDTreeCuda3dIndex {
+   public:
+    int visited_leafs;
+    KDTreeCuda3dIndex(const float* input_data, size_t input_count,
+                      const KDTreeCuda3dIndexParams& params = KDTreeCuda3dIndexParams())
+        : dataset_(input_data), leaf_count_(0), visited_leafs(0), node_count_(0), current_node_count_(0) {
+        size_ = input_count;
+        leaf_max_size_ = params.leaf_max_size;
+        gpu_helper_ = 0;
+    }
+
+    /**
+     * Standard destructor
+     */
+    ~KDTreeCuda3dIndex() { clearGpuBuffers(); }
+
+    /**
+     * Builds the index
+     */
+    void buildIndex() {
+        leaf_count_ = 0;
+        node_count_ = 0;
+        uploadTreeToGpu();
+    }
+
+    /**
+     * queries: (N, p) flattened float cuda array where only first 3 elements are used.
+     * n_query: N
+     * n_query_stride: p
+     * indices: (N, knn) int cuda array
+     * dists: (N, knn) float cuda array
+     */
+    void knnSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, size_t knn,
+                   const SearchParams& params = SearchParams()) const;
+    int radiusSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, float radius,
+                     const SearchParams& params = SearchParams()) const;
+
+   private:
+    void uploadTreeToGpu();
+    void clearGpuBuffers();
+
+   private:
+    struct GpuHelper;
+    GpuHelper* gpu_helper_;
+
+    const float* dataset_;
+    int leaf_max_size_;
+    int leaf_count_;
+    int node_count_;
+    //! used by convertTreeToGpuFormat
+    int current_node_count_;
+    size_t size_;
+
+};  // class KDTreeCuda3dIndex
+
+}  // namespace tinyflann
+
+#endif
@@ -0,0 +1,66 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAException.h>
+
+#include "cuda_kdtree.cuh"
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_IS_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, #x " must be a float tensor")
+
+std::vector<torch::Tensor> nearestNeighbours(torch::Tensor query, torch::Tensor tree, int knn) {
+    CHECK_CUDA(query);
+    CHECK_IS_FLOAT(query);
+    CHECK_CUDA(tree);
+    CHECK_IS_FLOAT(tree);
+    TORCH_CHECK(tree.size(0) >= knn, "knn is too small compared to the size of point cloud!");
+
+    torch::Tensor strided_tree = tree;
+    torch::Tensor strided_query = query;
+    torch::Device device = tree.device();
+    long n_queries = query.size(0);
+    long n_tree = tree.size(0);
+
+    if (strided_tree.stride(0) != 4) {
+        strided_tree = torch::zeros({n_tree, 4}, torch::dtype(torch::kFloat32).device(device));
+        strided_tree.index_put_(
+            {torch::indexing::Ellipsis, torch::indexing::Slice(torch::indexing::None, tree.size(1))}, tree);
+    }
+
+    if (strided_query.stride(0) != 3) {
+        strided_query = torch::zeros({n_queries, 3}, torch::dtype(torch::kFloat32).device(device));
+        strided_query.index_put_(
+            {torch::indexing::Ellipsis, torch::indexing::Slice(torch::indexing::None, query.size(1))}, query);
+    }
+
+    auto *knn_index = new tinyflann::KDTreeCuda3dIndex<tinyflann::CudaL2>(strided_tree.data_ptr<float>(), n_tree);
+    knn_index->buildIndex();
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    torch::Tensor dist = torch::empty(n_queries * knn, torch::dtype(torch::kFloat32).device(device));
+    torch::Tensor indices = torch::empty(n_queries * knn, torch::dtype(torch::kInt32).device(device));
+
+    knn_index->knnSearch(strided_query.data_ptr<float>(), n_queries, 3, indices.data_ptr<int>(), dist.data_ptr<float>(),
+                         knn);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    delete knn_index;
+
+    return {dist.reshape({n_queries, knn}), indices.reshape({n_queries, knn})};
+}
@@ -0,0 +1,50 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#define _CPU_AND_GPU_CODE_ __host__ __device__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CODE_
+#endif
+
+#if defined(__CUDACC__)
+#define _CPU_AND_GPU_CODE_TEMPLATE_ __host__ __device__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CODE_TEMPLATE_
+#endif
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#define _CPU_AND_GPU_CONSTANT_ __constant__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CONSTANT_ const
+#endif
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#include <THC/THCAtomics.cuh>
+
+template <typename ScalarT>
+static inline __device__ void atomicAdd(ScalarT* addr, ScalarT value) {
+    gpuAtomicAddNoReturn(addr, value);
+}
+#else
+template <typename ScalarT>
+static inline void atomicAdd(ScalarT* addr, ScalarT value) {
+    *addr += value;
+}
+#endif
@@ -0,0 +1,25 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+std::vector<torch::Tensor> nearestNeighbours(torch::Tensor query, torch::Tensor tree, int knn);
+
+void pybind_utils_ext(py::module &m) {
+    m.def("nearest_neighbours", &nearestNeighbours, "KNN computation (CUDA) distance is squared L2.", py::arg("query"),
+          py::arg("tree"), py::arg("knn"));
+}