mirror of
https://github.com/wassname/PointCNN.git
synced 2026-08-21 11:11:07 +08:00
83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
import unittest
|
|
|
|
import torch
|
|
from torch.autograd import Variable
|
|
import numpy as np
|
|
|
|
from PointCNN import XConv, RandPointCNN, knn_indices_func_cpu
|
|
from PointCNN.tests import plot_pts_and_fts
|
|
|
|
np.random.seed(0)
|
|
|
|
class BasicTests(unittest.TestCase):
|
|
""" Basic test cases """
|
|
|
|
def test_xconv_shape(self):
|
|
self.assertTrue(True)
|
|
|
|
N = 4
|
|
D = 3
|
|
C_in = 8
|
|
C_out = 32
|
|
N_neighbors = 100
|
|
|
|
model = XConv(C_in, C_out, D, N_neighbors).cuda()
|
|
p = Variable(torch.from_numpy(np.random.rand(N,D).astype(np.float32))).cuda()
|
|
P = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,D).astype(np.float32))).cuda()
|
|
F = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,C_in).astype(np.float32))).cuda()
|
|
out = model(p, P, F)
|
|
self.assertEqual(out.size(), (N, C_out))
|
|
|
|
def test_knn(self):
|
|
P = np.array([[[0,0],
|
|
[0,0.95],
|
|
[1,0],
|
|
[1,1]]])
|
|
ps = P[:,[0,3],:]
|
|
|
|
P = Variable(torch.from_numpy(P))
|
|
ps = Variable(torch.from_numpy(ps))
|
|
|
|
out = knn_indices_func_cpu(ps, P, 2).numpy()
|
|
target = np.array([[[1,2],
|
|
[2,1]]])
|
|
|
|
self.assertTrue(np.array_equal(target, out))
|
|
|
|
def test_pointcnn_shape(self):
|
|
N = 1
|
|
num_points = 1000
|
|
dims = 2
|
|
C_in = 4
|
|
K = 10
|
|
D = 1
|
|
|
|
layer1 = RandPointCNN(C_in, 8, dims, K, D, 1000, knn_indices_func_cpu).cuda()
|
|
layer2 = RandPointCNN( 8, 16, dims, K, D, 500, knn_indices_func_cpu).cuda()
|
|
layer3 = RandPointCNN( 16, 32, dims, K, D, 250, knn_indices_func_cpu).cuda()
|
|
layer4 = RandPointCNN( 32, 64, dims, K, D, 125, knn_indices_func_cpu).cuda()
|
|
layer5 = RandPointCNN( 64, 128, dims, K, D, 50, knn_indices_func_cpu).cuda()
|
|
|
|
pts = np.random.rand(N,num_points,dims).astype(np.float32)
|
|
fts = np.random.rand(N,num_points,C_in).astype(np.float32)
|
|
pts = Variable(torch.from_numpy(pts)).cuda()
|
|
fts = Variable(torch.from_numpy(fts)).cuda()
|
|
|
|
if True:
|
|
pts, fts = layer1((pts, fts))
|
|
else:
|
|
plot_pts_and_fts(pts, fts)
|
|
pts, fts = layer1((pts, fts))
|
|
plot_pts_and_fts(pts, fts)
|
|
pts, fts = layer2((pts, fts))
|
|
plot_pts_and_fts(pts, fts)
|
|
pts, fts = layer3((pts, fts))
|
|
plot_pts_and_fts(pts, fts)
|
|
pts, fts = layer4((pts, fts))
|
|
plot_pts_and_fts(pts, fts)
|
|
pts, fts = layer5((pts, fts))
|
|
plot_pts_and_fts(pts, fts)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|