knn interpolate dense

Author: nicolas-chaulet

.devcontainer/Dockerfile

@@ -30,3 +30,4 @@ RUN apt-get update \
 
 RUN pip3 install -U pip
 RUN pip3 install torch numpy scikit-learn flake8 setuptools
+RUN pip3 install torch_cluster torch_sparse torch_scatter torch_geometric

.devcontainer/devcontainer.json

@@ -4,27 +4,23 @@
 	"name": "Python 3",
 	"context": "..",
 	"dockerFile": "Dockerfile",
-
 	// Set *default* container specific settings.json values on container create.
-	"settings": { 
+	"settings": {
 		"terminal.integrated.shell.linux": "/bin/bash",
 		"python.pythonPath": "/usr/local/bin/python",
 		"python.linting.enabled": true,
 		"python.linting.pylintEnabled": true,
 		"python.linting.pylintPath": "/usr/local/bin/pylint"
 	},
-
 	// Add the IDs of extensions you want installed when the container is created.
 	"extensions": [
-		"ms-python.python"
+		"ms-python.python",
+		"ms-vscode.cpptools"
 	]
-
 	// Use 'forwardPorts' to make a list of ports inside the container available locally.
 	// "forwardPorts": [],
-
 	// Use 'postCreateCommand' to run commands after the container is created.
 	// "postCreateCommand": "pip install -r requirements.txt",
-
 	// Uncomment to connect as a non-root user. See https://aka.ms/vscode-remote/containers/non-root.
 	// "remoteUser": "vscode"
-}
+}

test/test_interpolate.py

@@ -0,0 +1,24 @@
+import unittest
+import torch
+from torch_points import three_interpolate_tg, three_interpolate, three_nn
+
+class TestInterpolate(unittest.TestCase):
+    def test_cpu(self):
+        pos = torch.randn([16, 100, 3])
+        pos_skip = torch.randn([16, 500, 3])
+        x = torch.randn([16, 30, 100])
+
+        # # dense
+        # dist, idx = three_nn(pos_skip, pos)
+        # dist_recip = 1.0 / (dist + 1e-8)
+        # norm = torch.sum(dist_recip, dim=2, keepdim=True)
+        # weight = dist_recip / norm
+        # interpolated_feats = three_interpolate(x, idx, weight)
+
+        # sparse
+        sp_interpolated = three_interpolate_tg(x,pos,pos_skip)
+
+        # torch.testing.assert_allclose(sp_interpolated, interpolated_feats)
+
+if __name__ == "__main__":
+    unittest.main()

torch_points/knn.py

@@ -11,4 +11,6 @@ def knn(pos_support, pos, k):
         idx - [B,M,k]
         dist2 - [B,M,k] squared distances
     """
+    assert pos_support.dim() == 3 and pos.dim() == 3
+    
     return tpcpu.dense_knn(pos_support, pos, k)

torch_points/knn_interpolate.py

@@ -0,0 +1,20 @@
+import torch
+from .knn import knn
+
+def knn_interpolate(x, pos, pos_support, k):
+    """ KNN interpolation for dense data
+
+    Parameters
+    ----------
+    x : (B, C, n) tensor of known features
+    pos : (B, n, 3) tensor of positions of known features
+    pos_support : (B, m, 3) tensor of position of unknown  features (generally m > n)
+
+    Returns
+    -------
+    (B, C, m) interpolated features
+    """
+    
+    knn_idx, knn_dist = knn(pos_support, pos, k)
+
+

torch_points/torchpoints.py

@@ -3,6 +3,7 @@
 import torch.nn as nn
 import sys
 from typing import Optional, Any, Tuple
+from torch_geometric.nn import knn_interpolate
 
 import torch_points.points_cpu as tpcpu
 
@@ -141,6 +142,13 @@ def three_interpolate(features, idx, weight):
     """
     return ThreeInterpolate.apply(features, idx, weight)
 
+def three_interpolate_tg(x, pos, new_pos):
+    interpolated = []
+    for i in range(x.shape[0]):
+        interpolated.append(knn_interpolate(x.transpose(1,0), pos, new_pos).transpose(1,0))
+    return torch.stack(interpolated)
+
+
 def grouping_operation(features, idx):
     r"""
     Parameters