-
Notifications
You must be signed in to change notification settings - Fork 5
/
generation.py
111 lines (88 loc) · 3.37 KB
/
generation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import torch
import torch.optim as optim
from torch import autograd
import numpy as np
import math
from tqdm import trange
import trimesh
import copy
import time
from tqdm import tqdm
from sklearn.neighbors import KDTree
import torch.nn.functional as F
import os
def rotation_matrix_from_vectors(vec1, vec2):
""" Find the rotation matrix that aligns vec1 to vec2
:param vec1: A 3d "source" vector
:param vec2: A 3d "destination" vector
:return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2.
"""
a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3)
v = np.cross(a, b)
if any(v): #if not all zeros then
c = np.dot(a, b)
s = np.linalg.norm(v)
kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
return np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2))
else:
return np.eye(3) #cross of all zeros only occurs on identical directions
class Generator3D6(object):
def __init__(self, model1, model2, device):
self.model1 = model1
self.model2 = model2
self.device = device
def upsample(self, data):
pc = self.generateiopoint(data)
return pc
def generateiopoint(self, data):
data=np.squeeze(data,0)
tree1=KDTree(data)
print("generate seedpoint")
wq="./dense"
wq=wq+" 0.004 "+str(data.shape[0])
print(wq)
os.system(wq)
data2=np.loadtxt("target.xyz")
xyz2=data2[:,0:3]
print("fn")
pp=xyz2.shape[0]//400
p_split = np.array_split(xyz2, pp, axis=0)
normal=None
for i in tqdm(range(len(p_split))):
dist, idx = tree1.query(p_split[i], 100)
cloud=data[idx]
cloud=cloud-np.tile(np.expand_dims(p_split[i],1),(1,100,1))
with torch.no_grad():
c = self.model1.encode_inputs(torch.from_numpy(np.expand_dims(cloud,0)).float().to(self.device))
with torch.no_grad():
n = self.model1.decode(c)
n=n.detach().cpu().numpy()
if normal is None:
normal=n
else:
normal= np.append(normal,n,axis=0)
n_split = np.array_split(normal, pp, axis=0)
xyzout=[]
print("fd")
for i in tqdm(range(len(p_split))):
dist, idx = tree1.query(p_split[i], 100)
cloud=data[idx]
cloud=cloud-np.tile(np.expand_dims(p_split[i],1),(1,100,1))
for j in range(cloud.shape[0]):
M1=rotation_matrix_from_vectors(n_split[i][j],[1,0,0])
cloud[j]=(np.matmul(M1,cloud[j].T)).T
with torch.no_grad():
c = self.model2.encode_inputs(torch.from_numpy(np.expand_dims(cloud,0)).float().to(self.device))
with torch.no_grad():
n = self.model2.decode(c)
length=np.tile(np.expand_dims(n.detach().cpu().numpy(),1),(1,3))
xyzout.extend((p_split[i]+n_split[i]*length).tolist())
xyzout=np.array(xyzout)
print("remove outliers")
tree3=KDTree(xyzout)
dist, idx = tree3.query(xyzout, 30)
avg=np.mean(dist,axis=1)
avgtotal=np.mean(dist)
idx=np.where(avg<avgtotal*1.5)[0]
xyzout=xyzout[idx,:]
return xyzout