pcn_emd.py
3.33 KB
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
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import tensorflow as tf
from tf_util import *
class Model:
def __init__(self, inputs,my_inputs, npts, gt, alpha,beta):
self.num_coarse = 1024
self.grid_size = 4
self.grid_scale = 0.05
self.num_fine = self.grid_size ** 2 * self.num_coarse
self.features = self.create_encoder(inputs, npts)
self.coarse, self.fine = self.create_decoder(self.features)
self.loss, self.update = self.create_loss(my_inputs,self.coarse, self.fine, gt, alpha,beta)
self.outputs = self.fine
self.visualize_ops = [tf.split(inputs[0], npts, axis=0), self.coarse, self.fine, gt]
self.visualize_titles = ['input', 'coarse output', 'fine output', 'ground truth']
def create_encoder(self, inputs, npts):
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
features = mlp_conv(inputs, [128, 256])
features_global = point_unpool(point_maxpool(features, npts, keepdims=True), npts)
features = tf.concat([features, features_global], axis=2)
with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE):
features = mlp_conv(features, [512, 1024])
features = point_maxpool(features, npts)
return features
def create_decoder(self, features):
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
coarse = mlp(features, [1024, 1024, self.num_coarse * 3])
coarse = tf.reshape(coarse, [-1, self.num_coarse, 3])
with tf.variable_scope('folding', reuse=tf.AUTO_REUSE):
x = tf.linspace(-self.grid_scale, self.grid_scale, self.grid_size)
y = tf.linspace(-self.grid_scale, self.grid_scale, self.grid_size)
grid = tf.meshgrid(x, y)
grid = tf.expand_dims(tf.reshape(tf.stack(grid, axis=2), [-1, 2]), 0)
grid_feat = tf.tile(grid, [features.shape[0], self.num_coarse, 1])
point_feat = tf.tile(tf.expand_dims(coarse, 2), [1, 1, self.grid_size ** 2, 1])
point_feat = tf.reshape(point_feat, [-1, self.num_fine, 3])
global_feat = tf.tile(tf.expand_dims(features, 1), [1, self.num_fine, 1])
feat = tf.concat([grid_feat, point_feat, global_feat], axis=2)
center = tf.tile(tf.expand_dims(coarse, 2), [1, 1, self.grid_size ** 2, 1])
center = tf.reshape(center, [-1, self.num_fine, 3])
fine = mlp_conv(feat, [512, 512, 3]) + center
return coarse, fine
def create_loss(self,my_inputs,coarse, fine, gt, alpha,beta):
gt_ds = gt[:, :coarse.shape[1], :]
loss_coarse = earth_mover(coarse, gt_ds)
add_train_summary('train/coarse_loss', loss_coarse)
update_coarse = add_valid_summary('valid/coarse_loss', loss_coarse)
loss_fine = chamfer(fine, gt)
add_train_summary('train/fine_loss', loss_fine)
update_fine = add_valid_summary('valid/fine_loss', loss_fine)
loss_input = my_chamfer(my_inputs,fine)
add_train_summary('train/input_loss',loss_input)
update_input = add_valid_summary('valid/input_loss',loss_input)
loss = loss_coarse + alpha * loss_fine + beta*loss_input
add_train_summary('train/loss', loss)
update_loss = add_valid_summary('valid/loss', loss)
return loss, [update_coarse, update_fine, update_input,update_loss]