Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
348 lines (323 sloc) 11.8 KB
import torch
from torch.utils import data
from scipy.io import loadmat
from enum import Enum
from torch.autograd import Variable
import math
from pytorch_ops.sampling.sample import FarthestSample
from pytorch_ops.losses.cd.cd import CDModule
m_grp = 0
def vrrotvec2mat(rotvector, angle):
s = math.sin(angle)
c = math.cos(angle)
t = 1 - c
x = rotvector[0]
y = rotvector[1]
z = rotvector[2]
m = torch.FloatTensor(
[[t * x * x + c, t * x * y - s * z, t * x * z + s * y],
[t * x * y + s * z, t * y * y + c, t * y * z - s * x],
[t * x * z - s * y, t * y * z + s * x, t * z * z + c]])
return m
#segmentation for symmetric node
def multilabel(points, shape, cdloss):
c = torch.LongTensor(1, 2048).zero_()
c = c - 1
for i in range(points.size(0)):
a = points[i].unsqueeze(0).cuda()
_, index, _, _ = cdloss(a, shape)
b = torch.unique(index.cpu())
for k in range(b.size(0)):
c[0, b[k].item()] = i
return c
class Tree(object):
class NodeType(Enum):
LEAF = 0 # leaf node
ADJ = 1 # adjacency (adjacent part assembly) node
SYM = 2 # symmetry (symmetric part grouping) node
SYM_ADJ = 3 #reflect
class Node(object):
def __init__(self,
leaf_points=None,
left=None,
right=None,
node_type=None,
sym_p=None,
sym_a=None,
sym_t=None,
semantic_label=None):
self.leaf_points = leaf_points # node points
if isinstance(sym_t, int):
self.sym_t = torch.LongTensor([sym_t])
else:
self.sym_t = None
if isinstance(sym_a, int):
self.sym_a = torch.LongTensor([sym_a])
else:
self.sym_a = None
self.sym_p = sym_p
self.sym_type = self.sym_a
self.left = left # left child for ADJ or SYM (a symmeter generator)
self.right = right # right child
self.node_type = node_type
self.label = torch.LongTensor([self.node_type.value])
self.is_root = False
self.semantic_label = semantic_label
def is_leaf(self):
return self.node_type == Tree.NodeType.LEAF
def is_adj(self):
return self.node_type == Tree.NodeType.ADJ
def is_sym(self):
return self.node_type == Tree.NodeType.SYM
def is_sym_adj(self):
return self.node_type == Tree.NodeType.SYM_ADJ
def __init__(self, parts, ops, syms, labels, shape):
parts_list = [p for p in torch.split(parts, 1, 0)]
sym_param = [s for s in torch.split(syms, 1, 0)]
part_labels = [s for s in torch.split(labels, 1, 0)]
parts_list.reverse()
sym_param.reverse()
part_labels.reverse()
queue = []
sym_node_num = 0
for id in range(ops.size()[1]):
if ops[0, id] == Tree.NodeType.LEAF.value:
queue.append(
Tree.Node(leaf_points=parts_list.pop(), node_type=Tree.NodeType.LEAF, semantic_label=part_labels.pop()))
elif ops[0, id] == Tree.NodeType.ADJ.value:
left_node = queue.pop()
right_node = queue.pop()
queue.append(
Tree.Node(
left=left_node,
right=right_node,
node_type=Tree.NodeType.ADJ))
elif ops[0, id] == Tree.NodeType.SYM.value:
node = queue.pop()
s = sym_param.pop()
b = s[0, 0] + 1
t = s[0, 7].item()
p = s[0, 1:7]
if t > 0:
t = round(1.0/t)
queue.append(
Tree.Node(
left=node,
sym_p=p.unsqueeze(0),
sym_a=int(b),
sym_t=int(t),
node_type=Tree.NodeType.SYM))
if b != 1:
sym_node_num += 1
assert len(queue) == 1
self.root = queue[0]
self.root.is_root = True
assert self.root.is_adj()
self.shape = shape
if sym_node_num == 0:
self.n_syms = torch.Tensor([sym_node_num]).cuda()
else:
self.n_syms = torch.Tensor([1/sym_node_num]).cuda()
#find GT label's index in input
def Attention(feature2048, shape):
index = []
for i in range(shape.size(1)):
if feature2048[0, i] > -1:
index.append(i)
pad_index = []
while len(pad_index) < 2048:
pad_index.extend(index)
pad_index = torch.LongTensor(pad_index[:2048])
return pad_index.unsqueeze(0).cpu()
#construct groundtruth for pointcloud segmentation
def dfs_fix(node, shape, cdloss, shape_normal, seg, grp, reflect=None):
global m_grp
if node.is_leaf():
# find node's corresponding points on input
_, index, _ , _ = cdloss(node.leaf_points[:, :, :3].cuda(), shape)
b = torch.unique(index.cpu())
c = torch.LongTensor(1, 2048).zero_()
c = c - 1
for i in range(b.size(0)):
c[0, b[i].item()] = 0
node.index = c #segmentation GT binary label
idx = Attention(c, shape) #node's corresponding idx
#node's corresponding points
node.points = torch.index_select(shape_normal, 1, idx.squeeze(0).long().cpu())
#node's corresponding idx
node.pad_index = idx
for i in range(node.pad_index.size(1)):
seg[node.pad_index[0, i].item()] = node.semantic_label
grp[node.pad_index[0, i].item()] = m_grp
m_grp += 1
if reflect is not None:
#recover reflect's children
re_leaf_points = torch.cat([node.leaf_points[:, :, :3], node.leaf_points[:, :, :3]+node.leaf_points[:, :, 3:]], 1)
re_leaf_points = re_leaf_points.squeeze(0).cpu()
sList = torch.split(reflect, 1, 0)
ref_normal = torch.cat([sList[0], sList[1], sList[2]])
ref_normal = ref_normal / torch.norm(ref_normal)
ref_point = torch.cat([sList[3], sList[4], sList[5]])
new_points = 2 * ref_point.add(-re_leaf_points).matmul(ref_normal)
new_points = new_points.unsqueeze(-1)
new_points = new_points.repeat(1, 3)
new_points = ref_normal.mul(new_points).add(re_leaf_points)
new_points = torch.cat([new_points[:2048, :], new_points[2048:, :] - new_points[:2048, :]], 1)
New_node = Tree.Node(leaf_points=new_points.unsqueeze(0), node_type=Tree.NodeType.LEAF)
#build node for reflect node's children
_, index, _ , _ = cdloss(New_node.leaf_points[:, :, :3].cuda(), shape)
b = torch.unique(index.cpu())
reflect_c = torch.LongTensor(1, 2048).zero_()
reflect_c = reflect_c - 1
for i in range(b.size(0)):
reflect_c[0, b[i].item()] = 0
New_node.index = reflect_c
idx = Attention(reflect_c, shape)
New_node.points = torch.index_select(shape_normal, 1, idx.squeeze(0).long().cpu())
New_node.pad_index = idx
New_node.semantic_label = node.semantic_label
for i in range(New_node.pad_index.size(1)):
seg[New_node.pad_index[0, i].item()] = New_node.semantic_label
grp[New_node.pad_index[0, i].item()] = m_grp
m_grp += 1
return torch.Tensor([0]).cuda(), New_node
else:
return torch.Tensor([0]).cuda(), node
if node.is_adj():
l_num, new_node_l = dfs_fix(node.left, shape, cdloss, shape_normal, seg, grp, reflect)
r_num, new_node_r = dfs_fix(node.right, shape, cdloss, shape_normal, seg, grp, reflect)
#build adj node
c = torch.LongTensor(1, 2048).zero_()
c = c - 1
for i in range(2048):
if node.left.index[0, i].item() > -1:
c[0, i] = 0
for i in range(2048):
if node.right.index[0, i].item() > -1:
c[0, i] = 1
node.index = c
idx = Attention(c, shape)
node.points = torch.index_select(shape_normal, 1, idx.squeeze(0).long().cpu())
node.pad_index = idx
if reflect is not None:
New_node = Tree.Node(left=new_node_l, right=new_node_r, node_type=Tree.NodeType.ADJ)
reflect_c = torch.LongTensor(1, 2048).zero_()
reflect_c = reflect_c - 1
for i in range(2048):
if new_node_l.index[0, i].item() > -1:
reflect_c[0, i] = 0
for i in range(2048):
if new_node_r.index[0, i].item() > -1:
reflect_c[0, i] = 1
New_node.index = reflect_c
idx = Attention(reflect_c, shape)
New_node.points = torch.index_select(shape_normal, 1, idx.squeeze(0).long().cpu())
New_node.pad_index = idx
return l_num + r_num + torch.Tensor([2]).cuda(), New_node
else:
return l_num + r_num + torch.Tensor([1]).cuda(), node
if node.is_sym():
#build symmetric node
t = node.sym_t.item()
p = node.sym_p.squeeze(0)
if node.sym_type.item() == 1: #reflect node
child_num, new_node = dfs_fix(node.left, shape, cdloss, shape_normal, seg, grp, p)
c = torch.LongTensor(1, 2048).zero_()
c = c - 1
for i in range(2048):
if node.left.index[0, i].item() > -1:
c[0, i] = 0
for i in range(2048):
if new_node.index[0, i].item() > -1:
c[0, i] = 1
node.index = c
node.right = new_node
idx = Attention(c, shape)
node.points = torch.index_select(shape_normal, 1, idx.squeeze(0).long().cpu())
node.node_type = Tree.NodeType.SYM_ADJ
node.label = torch.LongTensor([node.node_type.value])
node.pad_index = idx
return child_num + torch.Tensor([1]).cuda(), node
else:
child_num, _= dfs_fix(node.left, shape, cdloss, shape_normal, seg, grp, None)
new_leaf_points = node.left.leaf_points.squeeze(0)
leaf_points_list = [new_leaf_points.unsqueeze(0)]
new_leaf_points = torch.cat([new_leaf_points[:, :3] , new_leaf_points[:, :3] + new_leaf_points[:, 3:]], 0)
if node.sym_type.item() == 0:#rotate symmetry
sList = torch.split(p, 1, 0)
f1 = torch.cat([sList[0], sList[1], sList[2]])
if f1[1] < 0:
f1 = - f1
f1 = f1 / torch.norm(f1)
f2 = torch.cat([sList[3], sList[4], sList[5]])
folds = int(t)
a = 1.0 / float(folds)
for i in range(folds - 1):
angle = a * 2 * 3.1415 * (i + 1)
rotm = vrrotvec2mat(f1, angle)
sym_leaf_points = rotm.matmul(new_leaf_points.add(-f2).t()).t().add(f2)
sym_leaf_points = torch.cat([sym_leaf_points[:2048, :] , sym_leaf_points[2048:, :] - sym_leaf_points[:2048, :]], 1)
leaf_points_list.append(sym_leaf_points.unsqueeze(0))
elif node.sym_type.item() == 2: #translate symmetry
sList = torch.split(p, 1, 0)
trans = torch.cat([sList[0], sList[1], sList[2]])
folds = t - 1
trans = trans / float(folds)
for i in range(folds):
sym_leaf_points = new_leaf_points.add(trans.mul(i + 1))
sym_leaf_points = torch.cat([sym_leaf_points[:2048, :] , sym_leaf_points[2048:, :] - sym_leaf_points[:2048, :]], 1)
leaf_points_list.append(sym_leaf_points.unsqueeze(0))
a = torch.cat(leaf_points_list, 0)
node.index = multilabel(a[:, :, :3], shape, cdloss)
idx = Attention(node.index, shape)
node.points = torch.index_select(shape_normal, 1, idx.squeeze(0).long().cpu())
node.pad_index = Attention(node.index, shape)
for i in range(node.pad_index.size(1)):
seg[node.pad_index[0, i].item()] = node.left.semantic_label
for i in range(2048):
if node.index[0, i].item() > -1:
grp[i] = m_grp + node.index[0, i]
m_grp = m_grp + torch.max(node.index) + 1
return torch.Tensor([1]).cuda(), node
class Data_Loader(data.Dataset):
def __init__(self, dir, is_train, split_num, total_num):
self.dir = dir
op_data = torch.from_numpy(loadmat(self.dir + 'training_trees/ops.mat')['ops']).int()
label_data = torch.from_numpy(loadmat(self.dir + 'training_trees/labels.mat')['labels']).int()
sym_data = torch.from_numpy(loadmat(self.dir + 'training_trees/syms.mat')['syms']).float()
num_examples = op_data.size()[1]
op_data = torch.chunk(op_data, num_examples, 1)
label_data = torch.chunk(label_data, num_examples, 1)
sym_data = torch.chunk(sym_data, num_examples, 1)
self.trees = []
self.training = is_train
if is_train:
begin = 0
end = split_num
else:
begin = split_num
end = total_num
for i in range(begin, end):
parts = torch.from_numpy(loadmat(self.dir + 'training_data_models_segment_2048_normals/%d.mat' % i)['pc']).float()
shape = torch.from_numpy(loadmat(self.dir + 'models_2048_points_normals/%d.mat' % i)['pc']).float()
ops = torch.t(op_data[i])
syms = torch.t(sym_data[i])
labels = torch.t(label_data[i])
tree = Tree(parts, ops, syms, labels, shape)
cdloss = CDModule()
seg = torch.LongTensor(2048).zero_() # for ap calculation
grp = torch.LongTensor(2048).zero_()
global m_grp
m_grp = 0
num_node, _ = dfs_fix(tree.root, shape[0, :, :3].unsqueeze(0).cuda(), cdloss, shape, seg, grp)
tree.n_nodes = num_node
tree.shape_label = seg
tree.grp = grp
self.trees.append(tree)
print('load data', i)
print(len(self.trees))
def __getitem__(self, index):
tree = self.trees[index]
return tree
def __len__(self):
return len(self.trees)
You can’t perform that action at this time.