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,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})};
+}