├── README.md
├── code
    ├── README.txt
    ├── network.py
    ├── part_color_mapping.json
    ├── sphere.txt
    ├── test.py
    ├── train.py
    ├── train_results
    │   └── trained_models
    │   │   └── readme.txt
    └── utils
    │   └── tf_util.py
└── data
    ├── ShapeNet
        ├── test
        │   └── put_test_data_here.txt
        └── train
        │   └── put_train_data_here.txt
    └── put_your_data_here.txt


/README.md:
--------------------------------------------------------------------------------
 1 | # PointGrid: A Deep Network for 3D Shape Understanding
 2 | 
 3 | ## Prerequisites:
 4 | 1. Python (with necessary common libraries such as numpy, scipy, etc.)
 5 | 2. TensorFlow
 6 | 3. You need to prepare your data in *.mat file with the following format:
 7 | 	- 'points': N x 3 array (x, y, z coordinates of the point cloud)
 8 | 	- 'labels': N x 1 array (1-based integer per-point labels)
 9 | 	- 'category': scalar (0-based integer model category)
10 | 
11 | ## Train:
12 | 	python train.py
13 | ## Test:
14 | 	python test.py
15 | 
16 | If you find this code useful, please cite our work at <br />
17 | <pre>
18 | @article{PointGrid,
19 | 	author = {Truc Le and Ye Duan},
20 | 	titile = {{PointGrid: A Deep Network for 3D Shape Understanding}},
21 | 	journal = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
22 | 	month = {June},
23 | 	year = {2018},
24 | }
25 | </pre>
26 | 


--------------------------------------------------------------------------------
/code/README.txt:
--------------------------------------------------------------------------------
1 | # Train:
2 | 	python train.py
3 | 
4 | # Test:
5 | 	python test.py
6 | 


--------------------------------------------------------------------------------
/code/network.py:
--------------------------------------------------------------------------------
  1 | import tensorflow as tf
  2 | import tensorflow.contrib.slim as slim
  3 | import json
  4 | import numpy as np
  5 | from functools import partial
  6 | from scipy import stats
  7 | import math
  8 | import os
  9 | import sys
 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 11 | sys.path.append(os.path.dirname(BASE_DIR))
 12 | sys.path.append(os.path.join(BASE_DIR, './utils'))
 13 | import tf_util
 14 | 
 15 | N = 16 # grid size is N x N x N
 16 | K = 4 # each cell has K points
 17 | NUM_CATEGORY = 16
 18 | NUM_SEG_PART = 50+1
 19 | NUM_PER_POINT_FEATURES = 3
 20 | NUM_FEATURES = K * NUM_PER_POINT_FEATURES + 1
 21 | 
 22 | batch_norm = partial(slim.batch_norm, decay=0.9, scale=True, epsilon=1e-5, scope='bn', updates_collections=None)
 23 | 
 24 | def leak_relu(x, leak=0.1, scope=None):
 25 |     return tf.where(x >= 0, x, leak * x)
 26 | 
 27 | 
 28 | 
 29 | def integer_label_to_one_hot_label(integer_label):
 30 |     if (len(integer_label.shape) == 0):
 31 |         one_hot_label = np.zeros((NUM_CATEGORY))
 32 |         one_hot_label[integer_label] = 1
 33 |     elif (len(integer_label.shape) == 1):
 34 |         one_hot_label = np.zeros((integer_label.shape[0], NUM_SEG_PART))
 35 |         for i in range(integer_label.shape[0]):
 36 |             one_hot_label[i, integer_label[i]] = 1
 37 |     elif (len(integer_label.shape) == 4):
 38 |         one_hot_label = np.zeros((N, N, N, K, NUM_SEG_PART))
 39 |         for i in range(N):
 40 |           for j in range(N):
 41 |             for k in range(N):
 42 |               for l in range(K):
 43 |                 one_hot_label[i, j, k, l, integer_label[i, j, k, l]] = 1
 44 |     else:
 45 |         raise
 46 |     return one_hot_label
 47 | 
 48 | 
 49 | 
 50 | def pc2voxel(pc, pc_label):
 51 |     # Args:
 52 |     #     pc: size n x F where n is the number of points and F is feature size
 53 |     #     pc_label: size n x NUM_SEG_PART (one-hot encoding label)
 54 |     # Returns:
 55 |     #     voxel: N x N x N x K x (3+3)
 56 |     #     label: N x N x N x (K+1) x NUM_SEG_PART
 57 |     #     index: N x N x N x K
 58 | 
 59 |     num_points = pc.shape[0]
 60 |     data = np.zeros((N, N, N, NUM_FEATURES), dtype=np.float32)
 61 |     label = np.zeros((N, N, N, K+1, NUM_SEG_PART), dtype=np.float32)
 62 |     index = np.zeros((N, N, N, K), dtype=np.float32)
 63 |     xyz = pc[:, 0 : 3]
 64 |     centroid = np.mean(xyz, axis=0, keepdims=True)
 65 |     xyz -= centroid
 66 |     xyz /= np.amax(np.sqrt(np.sum(xyz ** 2, axis=1)), axis=0) * 1.05
 67 |     idx = np.floor((xyz + 1.0) / 2.0 * N)
 68 |     L = [[] for _ in range(N * N * N)]
 69 |     for p in range(num_points):
 70 |         k = int(idx[p, 0] * N * N + idx[p, 1] * N + idx[p, 2])
 71 |         L[k].append(p)
 72 |     for i in range(N):
 73 |       for j in range(N):
 74 |         for k in range(N):
 75 |           u = int(i * N * N + j * N + k)
 76 |           if not L[u]:
 77 |               data[i, j, k, :] = np.zeros((NUM_FEATURES), dtype=np.float32)
 78 |               label[i, j, k, :, :] = 0
 79 |               label[i, j, k, :, 0] = 1
 80 |           elif (len(L[u]) >= K):
 81 |               choice = np.random.choice(L[u], size=K, replace=False)
 82 |               local_points = pc[choice, :] - np.array([-1.0 + (i + 0.5) * 2.0 / N, -1.0 + (j + 0.5) * 2.0 / N, -1.0 + (k + 0.5) * 2.0 / N], dtype=np.float32)
 83 |               data[i, j, k, 0 : K * NUM_PER_POINT_FEATURES] = np.reshape(local_points, (K * NUM_PER_POINT_FEATURES))
 84 |               data[i, j, k, K * NUM_PER_POINT_FEATURES] = 1.0
 85 |               label[i, j, k, 0 : K, :] = pc_label[choice, :]
 86 |               majority = np.argmax(np.sum(pc_label[L[u], :], axis=0))
 87 |               label[i, j, k, K, :] = 0
 88 |               label[i, j, k, K, majority] = 1
 89 |               index[i, j, k, :] = choice
 90 |           else:
 91 |               choice = np.random.choice(L[u], size=K, replace=True)
 92 |               local_points = pc[choice, :] - np.array([-1.0 + (i + 0.5) * 2.0 / N, -1.0 + (j + 0.5) * 2.0 / N, -1.0 + (k + 0.5) * 2.0 / N], dtype=np.float32)
 93 |               data[i, j, k, 0 : K * NUM_PER_POINT_FEATURES] = np.reshape(local_points, (K * NUM_PER_POINT_FEATURES))
 94 |               data[i, j, k, K * NUM_PER_POINT_FEATURES] = 1.0
 95 |               label[i, j, k, 0 : K, :] = pc_label[choice, :]
 96 |               majority = np.argmax(np.sum(pc_label[L[u], :], axis=0))
 97 |               label[i, j, k, K, :] = 0
 98 |               label[i, j, k, K, majority] = 1
 99 |               index[i, j, k, :] = choice
100 |     return data, label, index
101 | 
102 | 
103 | 
104 | def rotate_pc(pc):
105 |     # Args:
106 |     #     pc: size n x 3
107 |     # Returns:
108 |     #     rotated_pc: size n x 3
109 |     angle = np.random.uniform() * 2 * np.pi
110 |     cosval = np.cos(angle)
111 |     sinval = np.sin(angle)
112 |     rotation_matrix = np.array([[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]])
113 |     rotated_pc = np.dot(pc, rotation_matrix)
114 |     return rotated_pc
115 | 
116 | 
117 | 
118 | def populateIntegerSegLabel(pc, voxel_label, index):
119 |     # Args:
120 |     #     pc: size n x F where n is the number of points and F is feature size
121 |     #     voxel_label: size N x N x N x (K+1)
122 |     #     index: size N x N x N x K
123 |     # Returns:
124 |     #     label: size n x 1
125 | 
126 |     num_points = pc.shape[0]
127 |     label = np.zeros((num_points), dtype=np.int32)
128 |     xyz = pc[:, 0 : 3]
129 |     centroid = np.mean(xyz, axis=0, keepdims=True)
130 |     xyz -= centroid
131 |     xyz /= np.amax(np.sqrt(np.sum(xyz ** 2, axis=1)), axis=0) * 1.05
132 |     idx = np.floor((xyz + 1.0) / 2.0 * N)
133 |     L = [[] for _ in range(N * N * N)]
134 |     for p in range(num_points):
135 |         k = int(idx[p, 0] * N * N + idx[p, 1] * N + idx[p, 2])
136 |         L[k].append(p)
137 |     for i in range(N):
138 |       for j in range(N):
139 |         for k in range(N):
140 |           u = int(i * N * N + j * N + k)
141 |           if not L[u]:
142 |               pass
143 |           elif (len(L[u]) >= K):
144 |               label[L[u]] = voxel_label[i, j, k, K]
145 |               for s in range(K):
146 |                   label[int(index[int(i), int(j), int(k), int(s)])] = int(voxel_label[int(i), int(j), int(k), int(s)])
147 |           else:
148 |               for s in range(K):
149 |                   label[int(index[int(i), int(j), int(k), int(s)])] = int(voxel_label[int(i), int(j), int(k), int(s)])
150 |     return label
151 | 
152 | 
153 | 
154 | def populateOneHotSegLabel(pc, voxel_label, index):
155 |     # Args:
156 |     #     pc: size n x F where n is the number of points and F is feature size
157 |     #     voxel_label: size N x N x N x (K+1) x NUM_SEG_PART
158 |     #     index: size N x N x N x K
159 |     # Returns:
160 |     #     label: size n x 1
161 |     return populateIntegerSegLabel(pc, np.argmax(voxel_label, axis=4), index)
162 | 
163 | 
164 | 
165 | def get_model(pointgrid, is_training):
166 |     # Args:
167 |     #     pointgrid: of size B x N x N x N x NUM_FEATURES
168 |     #     is_training: boolean tensor
169 |     # Returns:
170 |     #     pred_cat: of size B x NUM_CATEGORY
171 |     #     pred_seg: of size B x N x N x N x (K+1) x NUM_PART_SEG
172 | 
173 |     # Encoder
174 |     batch_size = pointgrid.get_shape()[0].value
175 |     conv1 = tf_util.conv3d(pointgrid, 64, [5,5,5], scope='conv1', activation_fn=leak_relu, bn=True, is_training=is_training) # N
176 |     conv2 = tf_util.conv3d(conv1, 64, [5,5,5], scope='conv2', activation_fn=leak_relu, stride=[2,2,2], bn=True, is_training=is_training) # N/2
177 |     conv3 = tf_util.conv3d(conv2, 64, [5,5,5], scope='conv3', activation_fn=leak_relu, bn=True, is_training=is_training) # N/2
178 |     conv4 = tf_util.conv3d(conv3, 128, [3,3,3], scope='conv4', activation_fn=leak_relu, stride=[2,2,2], bn=True, is_training=is_training) # N/4
179 |     conv5 = tf_util.conv3d(conv4, 128, [3,3,3], scope='conv5', activation_fn=leak_relu, bn=True, is_training=is_training) # N/4
180 |     conv6 = tf_util.conv3d(conv5, 256, [3,3,3], scope='conv6', activation_fn=leak_relu, stride=[2,2,2], bn=True, is_training=is_training) # N/8
181 |     conv7 = tf_util.conv3d(conv6, 256, [3,3,3], scope='conv7', activation_fn=leak_relu, bn=True, is_training=is_training) # N/8
182 |     conv8 = tf_util.conv3d(conv7, 512, [3,3,3], scope='conv8', activation_fn=leak_relu, stride=[2,2,2], bn=True, is_training=is_training) # N/16
183 |     conv9 = tf_util.conv3d(conv8, 512, [1,1,1], scope='conv9', activation_fn=leak_relu, bn=True, is_training=is_training) # N/16
184 | 
185 |     # Classification Network
186 |     conv9_flat = tf.reshape(conv9, [batch_size, -1])
187 |     fc1 = tf_util.fully_connected(conv9_flat, 512, activation_fn=leak_relu, bn=True, is_training=is_training, scope='fc1')
188 |     do1 = tf_util.dropout(fc1, keep_prob=0.7, is_training=is_training, scope='do1')
189 |     fc2 = tf_util.fully_connected(do1, 256, activation_fn=leak_relu, bn=True, is_training=is_training, scope='fc2')
190 |     do2 = tf_util.dropout(fc2, keep_prob=0.7, is_training=is_training, scope='do2')
191 |     pred_cat = tf_util.fully_connected(do2, NUM_CATEGORY, activation_fn=None, bn=False, scope='pred_cat')
192 | 
193 |     # Segmentation Network
194 |     cat_features = tf.tile(tf.reshape(tf.concat([fc2, pred_cat], axis=1), [batch_size, 1, 1, 1, -1]), [1, N/16, N/16, N/16, 1])
195 |     conv9_cat = tf.concat([conv9, cat_features], axis=4)
196 |     deconv1 = tf_util.conv3d_transpose(conv9_cat, 256, [3,3,3], scope='deconv1', activation_fn=leak_relu, bn=True, is_training=is_training, stride=[2,2,2], padding='SAME') # N/8
197 |     conv7_deconv1 = tf.concat(axis=4, values=[conv7, deconv1])
198 |     deconv2 = tf_util.conv3d(conv7_deconv1, 256, [3,3,3], scope='deconv2', activation_fn=leak_relu, bn=True, is_training=is_training) # N/8
199 |     deconv3 = tf_util.conv3d_transpose(deconv2, 128, [3,3,3], scope='deconv3', activation_fn=leak_relu, bn=True, is_training=is_training, stride=[2,2,2], padding='SAME') # N/4
200 |     conv5_deconv3 = tf.concat(axis=4, values=[conv5, deconv3])
201 |     deconv4 = tf_util.conv3d(conv5_deconv3, 128, [3,3,3], scope='deconv4', activation_fn=leak_relu, bn=True, is_training=is_training) # N/4
202 |     deconv5 = tf_util.conv3d_transpose(deconv4, 64, [3,3,3], scope='deconv5', activation_fn=leak_relu, bn=True, is_training=is_training, stride=[2,2,2], padding='SAME') # N/2
203 |     conv3_deconv5 = tf.concat(axis=4, values=[conv3, deconv5])
204 |     deconv6 = tf_util.conv3d(conv3_deconv5, 64, [5,5,5], scope='deconv6', activation_fn=leak_relu, bn=True, is_training=is_training) # N/2
205 |     deconv7 = tf_util.conv3d_transpose(deconv6, 64, [5,5,5], scope='deconv7', activation_fn=leak_relu, bn=True, is_training=is_training, stride=[2,2,2], padding='SAME') # N
206 |     conv1_deconv7 = tf.concat(axis=4, values=[conv1, deconv7])
207 |     deconv8 = tf_util.conv3d(conv1_deconv7, 64, [5,5,5], scope='deconv8', activation_fn=leak_relu, bn=True, is_training=is_training) # N
208 | 
209 |     pred_seg = tf_util.conv3d(deconv8, (K+1) * NUM_SEG_PART, [5,5,5], scope='pred_seg', activation_fn=None, bn=False, is_training=is_training)
210 |     pred_seg = tf.reshape(pred_seg, [batch_size, N, N, N, K+1, NUM_SEG_PART])
211 | 
212 |     return pred_cat, pred_seg
213 | 
214 | 
215 | 
216 | def get_loss(pred_cat, one_hot_cat, pred_seg, one_hot_seg):
217 |     per_instance_cat_loss = tf.nn.softmax_cross_entropy_with_logits(logits=pred_cat, labels=one_hot_cat)
218 |     cat_loss = tf.constant(200.0, dtype=tf.float32) * tf.reduce_mean(per_instance_cat_loss)
219 |     per_instance_seg_loss = tf.nn.softmax_cross_entropy_with_logits(logits=pred_seg, labels=one_hot_seg)
220 |     seg_loss = tf.constant(800.0, dtype=tf.float32) * tf.reduce_mean(per_instance_seg_loss)
221 |     total_var = tf.trainable_variables()
222 |     reg_vars = [var for var in total_var if 'weights' in var.name]
223 |     reg_loss = tf.zeros([], dtype=tf.float32)
224 |     for var in reg_vars:
225 |         reg_loss += tf.nn.l2_loss(var)
226 |     reg_loss = tf.constant(1e-5, dtype=tf.float32) * reg_loss
227 |     total_loss = cat_loss + seg_loss + reg_loss
228 |     return total_loss, cat_loss, seg_loss
229 | 
230 | 
231 | 
232 | def intersection_over_union(pred_seg, integer_seg_label):
233 |     iou, counts = 0.0, 0.0
234 |     for i in range(1, NUM_SEG_PART):
235 |         intersection = np.sum(np.logical_and(pred_seg == i, integer_seg_label == i))
236 |         union = np.sum(np.logical_or(pred_seg == i, integer_seg_label == i))
237 |         if (union > 0):
238 |             counts += 1.0
239 |             iou += (float(intersection) / float(union))
240 |     iou /= counts
241 |     return iou
242 | 


--------------------------------------------------------------------------------
/code/part_color_mapping.json:
--------------------------------------------------------------------------------
1 | [[0.65, 0.95, 0.05], [0.35, 0.05, 0.35], [0.65, 0.35, 0.65], [0.95, 0.95, 0.65], [0.95, 0.65, 0.05], [0.35, 0.05, 0.05], [0.65, 0.05, 0.05], [0.65, 0.35, 0.95], [0.05, 0.05, 0.65], [0.65, 0.05, 0.35], [0.05, 0.35, 0.35], [0.65, 0.65, 0.35], [0.35, 0.95, 0.05], [0.05, 0.35, 0.65], [0.95, 0.95, 0.35], [0.65, 0.65, 0.65], [0.95, 0.95, 0.05], [0.65, 0.35, 0.05], [0.35, 0.65, 0.05], [0.95, 0.65, 0.95], [0.95, 0.35, 0.65], [0.05, 0.65, 0.95], [0.65, 0.95, 0.65], [0.95, 0.35, 0.95], [0.05, 0.05, 0.95], [0.65, 0.05, 0.95], [0.65, 0.05, 0.65], [0.35, 0.35, 0.95], [0.95, 0.95, 0.95], [0.05, 0.05, 0.05], [0.05, 0.35, 0.95], [0.65, 0.95, 0.95], [0.95, 0.05, 0.05], [0.35, 0.95, 0.35], [0.05, 0.35, 0.05], [0.05, 0.65, 0.35], [0.05, 0.95, 0.05], [0.95, 0.65, 0.65], [0.35, 0.95, 0.95], [0.05, 0.95, 0.35], [0.95, 0.35, 0.05], [0.65, 0.35, 0.35], [0.35, 0.95, 0.65], [0.35, 0.35, 0.65], [0.65, 0.95, 0.35], [0.05, 0.95, 0.65], [0.65, 0.65, 0.95], [0.35, 0.05, 0.95], [0.35, 0.65, 0.95], [0.35, 0.05, 0.65]]
2 | 


--------------------------------------------------------------------------------
/code/sphere.txt:
--------------------------------------------------------------------------------
  1 | 114
  2 | 0.38268343 0.92387953 0.0000000e+0
  3 | 0.35355339 0.92387953 0.14644661
  4 | 0.27059805 0.92387953 0.27059805
  5 | 0.14644661 0.92387953 0.35355339
  6 | 2.3432602e-17 0.92387953 0.38268343
  7 | -0.14644661 0.92387953 0.35355339
  8 | -0.27059805 0.92387953 0.27059805
  9 | -0.35355339 0.92387953 0.14644661
 10 | -0.38268343 0.92387953 4.6865204e-17
 11 | -0.35355339 0.92387953 -0.14644661
 12 | -0.27059805 0.92387953 -0.27059805
 13 | -0.14644661 0.92387953 -0.35355339
 14 | -7.0297806e-17 0.92387953 -0.38268343
 15 | 0.14644661 0.92387953 -0.35355339
 16 | 0.27059805 0.92387953 -0.27059805
 17 | 0.35355339 0.92387953 -0.14644661
 18 | 0.70710678 0.70710678 0.0000000e+0
 19 | 0.65328148 0.70710678 0.27059805
 20 | 0.50000000 0.70710678 0.50000000
 21 | 0.27059805 0.70710678 0.65328148
 22 | 4.3297803e-17 0.70710678 0.70710678
 23 | -0.27059805 0.70710678 0.65328148
 24 | -0.50000000 0.70710678 0.50000000
 25 | -0.65328148 0.70710678 0.27059805
 26 | -0.70710678 0.70710678 8.6595606e-17
 27 | -0.65328148 0.70710678 -0.27059805
 28 | -0.50000000 0.70710678 -0.50000000
 29 | -0.27059805 0.70710678 -0.65328148
 30 | -1.2989341e-16 0.70710678 -0.70710678
 31 | 0.27059805 0.70710678 -0.65328148
 32 | 0.50000000 0.70710678 -0.50000000
 33 | 0.65328148 0.70710678 -0.27059805
 34 | 0.92387953 0.38268343 0.0000000e+0
 35 | 0.85355339 0.38268343 0.35355339
 36 | 0.65328148 0.38268343 0.65328148
 37 | 0.35355339 0.38268343 0.85355339
 38 | 5.6571306e-17 0.38268343 0.92387953
 39 | -0.35355339 0.38268343 0.85355339
 40 | -0.65328148 0.38268343 0.65328148
 41 | -0.85355339 0.38268343 0.35355339
 42 | -0.92387953 0.38268343 1.1314261e-16
 43 | -0.85355339 0.38268343 -0.35355339
 44 | -0.65328148 0.38268343 -0.65328148
 45 | -0.35355339 0.38268343 -0.85355339
 46 | -1.6971392e-16 0.38268343 -0.92387953
 47 | 0.35355339 0.38268343 -0.85355339
 48 | 0.65328148 0.38268343 -0.65328148
 49 | 0.85355339 0.38268343 -0.35355339
 50 | 1.00000000 6.1232340e-17 0.0000000e+0
 51 | 0.92387953 6.1232340e-17 0.38268343
 52 | 0.70710678 6.1232340e-17 0.70710678
 53 | 0.38268343 6.1232340e-17 0.92387953
 54 | 6.1232340e-17 6.1232340e-17 1.00000000
 55 | -0.38268343 6.1232340e-17 0.92387953
 56 | -0.70710678 6.1232340e-17 0.70710678
 57 | -0.92387953 6.1232340e-17 0.38268343
 58 | -1.00000000 6.1232340e-17 1.2246468e-16
 59 | -0.92387953 6.1232340e-17 -0.38268343
 60 | -0.70710678 6.1232340e-17 -0.70710678
 61 | -0.38268343 6.1232340e-17 -0.92387953
 62 | -1.8369702e-16 6.1232340e-17 -1.00000000
 63 | 0.38268343 6.1232340e-17 -0.92387953
 64 | 0.70710678 6.1232340e-17 -0.70710678
 65 | 0.92387953 6.1232340e-17 -0.38268343
 66 | 0.92387953 -0.38268343 0.0000000e+0
 67 | 0.85355339 -0.38268343 0.35355339
 68 | 0.65328148 -0.38268343 0.65328148
 69 | 0.35355339 -0.38268343 0.85355339
 70 | 5.6571306e-17 -0.38268343 0.92387953
 71 | -0.35355339 -0.38268343 0.85355339
 72 | -0.65328148 -0.38268343 0.65328148
 73 | -0.85355339 -0.38268343 0.35355339
 74 | -0.92387953 -0.38268343 1.1314261e-16
 75 | -0.85355339 -0.38268343 -0.35355339
 76 | -0.65328148 -0.38268343 -0.65328148
 77 | -0.35355339 -0.38268343 -0.85355339
 78 | -1.6971392e-16 -0.38268343 -0.92387953
 79 | 0.35355339 -0.38268343 -0.85355339
 80 | 0.65328148 -0.38268343 -0.65328148
 81 | 0.85355339 -0.38268343 -0.35355339
 82 | 0.70710678 -0.70710678 0.0000000e+0
 83 | 0.65328148 -0.70710678 0.27059805
 84 | 0.50000000 -0.70710678 0.50000000
 85 | 0.27059805 -0.70710678 0.65328148
 86 | 4.3297803e-17 -0.70710678 0.70710678
 87 | -0.27059805 -0.70710678 0.65328148
 88 | -0.50000000 -0.70710678 0.50000000
 89 | -0.65328148 -0.70710678 0.27059805
 90 | -0.70710678 -0.70710678 8.6595606e-17
 91 | -0.65328148 -0.70710678 -0.27059805
 92 | -0.50000000 -0.70710678 -0.50000000
 93 | -0.27059805 -0.70710678 -0.65328148
 94 | -1.2989341e-16 -0.70710678 -0.70710678
 95 | 0.27059805 -0.70710678 -0.65328148
 96 | 0.50000000 -0.70710678 -0.50000000
 97 | 0.65328148 -0.70710678 -0.27059805
 98 | 0.38268343 -0.92387953 0.0000000e+0
 99 | 0.35355339 -0.92387953 0.14644661
100 | 0.27059805 -0.92387953 0.27059805
101 | 0.14644661 -0.92387953 0.35355339
102 | 2.3432602e-17 -0.92387953 0.38268343
103 | -0.14644661 -0.92387953 0.35355339
104 | -0.27059805 -0.92387953 0.27059805
105 | -0.35355339 -0.92387953 0.14644661
106 | -0.38268343 -0.92387953 4.6865204e-17
107 | -0.35355339 -0.92387953 -0.14644661
108 | -0.27059805 -0.92387953 -0.27059805
109 | -0.14644661 -0.92387953 -0.35355339
110 | -7.0297806e-17 -0.92387953 -0.38268343
111 | 0.14644661 -0.92387953 -0.35355339
112 | 0.27059805 -0.92387953 -0.27059805
113 | 0.35355339 -0.92387953 -0.14644661
114 | 0.0000000e+0 1.00000000 0.0000000e+0
115 | 0.0000000e+0 -1.00000000 0.0000000e+0
116 | 224
117 | 1 18 17
118 | 1 32 16
119 | 1 113 2
120 | 2 18 1
121 | 2 113 3
122 | 3 18 2
123 | 3 20 19
124 | 3 113 4
125 | 4 20 3
126 | 4 113 5
127 | 5 20 4
128 | 5 22 21
129 | 5 113 6
130 | 6 22 5
131 | 6 113 7
132 | 7 22 6
133 | 7 24 23
134 | 7 113 8
135 | 8 24 7
136 | 8 113 9
137 | 9 24 8
138 | 9 26 25
139 | 9 113 10
140 | 10 26 9
141 | 10 113 11
142 | 11 26 10
143 | 11 28 27
144 | 11 113 12
145 | 12 28 11
146 | 12 113 13
147 | 13 28 12
148 | 13 30 29
149 | 13 113 14
150 | 14 30 13
151 | 14 113 15
152 | 15 30 14
153 | 15 32 31
154 | 15 113 16
155 | 16 32 15
156 | 16 113 1
157 | 17 32 1
158 | 17 33 32
159 | 18 33 17
160 | 18 35 34
161 | 19 18 3
162 | 19 35 18
163 | 20 35 19
164 | 20 37 36
165 | 21 20 5
166 | 21 37 20
167 | 22 37 21
168 | 22 39 38
169 | 23 22 7
170 | 23 39 22
171 | 24 39 23
172 | 24 41 40
173 | 25 24 9
174 | 25 41 24
175 | 26 41 25
176 | 26 43 42
177 | 27 26 11
178 | 27 43 26
179 | 28 43 27
180 | 28 45 44
181 | 29 28 13
182 | 29 45 28
183 | 30 45 29
184 | 30 47 46
185 | 31 30 15
186 | 31 47 30
187 | 32 33 48
188 | 32 47 31
189 | 33 50 49
190 | 33 64 48
191 | 34 33 18
192 | 34 50 33
193 | 35 50 34
194 | 35 52 51
195 | 36 35 20
196 | 36 52 35
197 | 37 52 36
198 | 37 54 53
199 | 38 37 22
200 | 38 54 37
201 | 39 54 38
202 | 39 56 55
203 | 40 39 24
204 | 40 56 39
205 | 41 56 40
206 | 41 58 57
207 | 42 41 26
208 | 42 58 41
209 | 43 58 42
210 | 43 60 59
211 | 44 43 28
212 | 44 60 43
213 | 45 60 44
214 | 45 62 61
215 | 46 45 30
216 | 46 62 45
217 | 47 62 46
218 | 47 64 63
219 | 48 47 32
220 | 48 64 47
221 | 49 64 33
222 | 49 65 64
223 | 50 65 49
224 | 50 67 66
225 | 51 50 35
226 | 51 67 50
227 | 52 67 51
228 | 52 69 68
229 | 53 52 37
230 | 53 69 52
231 | 54 69 53
232 | 54 71 70
233 | 55 54 39
234 | 55 71 54
235 | 56 71 55
236 | 56 73 72
237 | 57 56 41
238 | 57 73 56
239 | 58 73 57
240 | 58 75 74
241 | 59 58 43
242 | 59 75 58
243 | 60 75 59
244 | 60 77 76
245 | 61 60 45
246 | 61 77 60
247 | 62 77 61
248 | 62 79 78
249 | 63 62 47
250 | 63 79 62
251 | 64 65 80
252 | 64 79 63
253 | 65 82 81
254 | 65 96 80
255 | 66 65 50
256 | 66 82 65
257 | 67 82 66
258 | 67 84 83
259 | 68 67 52
260 | 68 84 67
261 | 69 84 68
262 | 69 86 85
263 | 70 69 54
264 | 70 86 69
265 | 71 86 70
266 | 71 88 87
267 | 72 71 56
268 | 72 88 71
269 | 73 88 72
270 | 73 90 89
271 | 74 73 58
272 | 74 90 73
273 | 75 90 74
274 | 75 92 91
275 | 76 75 60
276 | 76 92 75
277 | 77 92 76
278 | 77 94 93
279 | 78 77 62
280 | 78 94 77
281 | 79 94 78
282 | 79 96 95
283 | 80 79 64
284 | 80 96 79
285 | 81 96 65
286 | 81 97 96
287 | 82 97 81
288 | 82 99 98
289 | 83 82 67
290 | 83 99 82
291 | 84 99 83
292 | 84 101 100
293 | 85 84 69
294 | 85 101 84
295 | 86 101 85
296 | 86 103 102
297 | 87 86 71
298 | 87 103 86
299 | 88 103 87
300 | 88 105 104
301 | 89 88 73
302 | 89 105 88
303 | 90 105 89
304 | 90 107 106
305 | 91 90 75
306 | 91 107 90
307 | 92 107 91
308 | 92 109 108
309 | 93 92 77
310 | 93 109 92
311 | 94 109 93
312 | 94 111 110
313 | 95 94 79
314 | 95 111 94
315 | 96 97 112
316 | 96 111 95
317 | 97 114 112
318 | 98 97 82
319 | 98 114 97
320 | 99 114 98
321 | 100 99 84
322 | 100 114 99
323 | 101 114 100
324 | 102 101 86
325 | 102 114 101
326 | 103 114 102
327 | 104 103 88
328 | 104 114 103
329 | 105 114 104
330 | 106 105 90
331 | 106 114 105
332 | 107 114 106
333 | 108 107 92
334 | 108 114 107
335 | 109 114 108
336 | 110 109 94
337 | 110 114 109
338 | 111 114 110
339 | 112 111 96
340 | 112 114 111


--------------------------------------------------------------------------------
/code/test.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import tensorflow as tf
  3 | import json
  4 | import numpy as np
  5 | import scipy.io
  6 | import os
  7 | import sys
  8 | import glob
  9 | from skimage import measure
 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 11 | sys.path.append(BASE_DIR)
 12 | sys.path.append(os.path.dirname(BASE_DIR))
 13 | import network as model
 14 | 
 15 | parser = argparse.ArgumentParser()
 16 | FLAGS = parser.parse_args()
 17 | 
 18 | 
 19 | # DEFAULT SETTINGS
 20 | gpu_to_use = 0
 21 | output_dir = os.path.join(BASE_DIR, './test_results')
 22 | 
 23 | # MAIN SCRIPT
 24 | batch_size = 1               # DO NOT CHANGE
 25 | purify = True                # Reassign label based on k-nearest neighbor. Set to False for large point cloud due to slow speed
 26 | knn = 5                      # for the purify
 27 | 
 28 | def get_file_name(file_path):
 29 |     parts = file_path.split('/')
 30 |     part = parts[-1]
 31 |     parts = part.split('.')
 32 |     return parts[0]
 33 | 
 34 | TESTING_FILE_LIST = [get_file_name(file_name) for file_name in glob.glob('../data/ShapeNet/test/' + '*.mat')]
 35 | 
 36 | category2name = ['Airplane', 'Bag', 'Cap', 'Car', 'Chair', 'Earphone', 'Guitar', 'Knife', 'Lamp', 'Laptop', 'Motorbike', 'Mug', 'Pistol', 'Rocket', 'Skateboard', 'Table']
 37 | 
 38 | color_map = json.load(open('part_color_mapping.json', 'r'))
 39 | 
 40 | lines = [line.rstrip('\n') for line in open('sphere.txt')]
 41 | nSphereVertices = int(lines[0])
 42 | sphereVertices = np.zeros((nSphereVertices, 3))
 43 | for i in range(nSphereVertices):
 44 |     coordinates = lines[i + 1].split()
 45 |     for j in range(len(coordinates)):
 46 |         sphereVertices[i, j] = float(coordinates[j])
 47 | nSphereFaces = int(lines[nSphereVertices + 1])
 48 | sphereFaces = np.zeros((nSphereFaces, 3))
 49 | for i in range(nSphereFaces):
 50 |     indices = lines[i + nSphereVertices + 2].split()
 51 |     for j in range(len(coordinates)):
 52 |         sphereFaces[i, j] = int(indices[j])
 53 | 
 54 | def output_color_point_cloud(data, seg, out_file, r=0.01):
 55 |     count = 0
 56 |     with open(out_file, 'w') as f:
 57 |         l = len(seg)
 58 |         for i in range(l):
 59 |             color = color_map[seg[i]]
 60 |             for j in range(nSphereVertices):
 61 |                 f.write('v %f %f %f %f %f %f\n' % \
 62 |                         (data[i][0] + sphereVertices[j][0] * r, data[i][1] + sphereVertices[j][1] * r, data[i][2] + sphereVertices[j][2] * r, color[0], color[1], color[2]))
 63 |             for j in range(nSphereFaces):
 64 |                 f.write('f %d %d %d\n' % (count + sphereFaces[j][0], count + sphereFaces[j][1], count + sphereFaces[j][2]))
 65 |             count += nSphereVertices
 66 | 
 67 | def printout(flog, data):
 68 |     print(data)
 69 |     flog.write(data + '\n')
 70 | 
 71 | def placeholder_inputs():
 72 |     pointgrid_ph = tf.placeholder(tf.float32, shape=(batch_size, model.N, model.N, model.N, model.NUM_FEATURES))
 73 |     cat_label_ph = tf.placeholder(tf.float32, shape=(batch_size, model.NUM_CATEGORY))
 74 |     seg_label_ph = tf.placeholder(tf.float32, shape=(batch_size, model.N, model.N, model.N, model.K+1, model.NUM_SEG_PART))
 75 |     return pointgrid_ph, cat_label_ph, seg_label_ph
 76 | 
 77 | def sigmoid(x):
 78 |     return 1 / (1 + np.exp(-x))
 79 | 
 80 | def load_checkpoint(checkpoint_dir, session, var_list=None):
 81 |     print(' [*] Loading checkpoint...')
 82 |     ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
 83 |     if ckpt and ckpt.model_checkpoint_path:
 84 |         ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
 85 |         ckpt_path = os.path.join(checkpoint_dir, ckpt_name)
 86 |     try:
 87 |         restorer = tf.train.Saver(var_list)
 88 |         restorer.restore(session, ckpt_path)
 89 |         print(' [*] Loading successful! Copy variables from % s' % ckpt_path)
 90 |         return True
 91 |     except:
 92 |         print(' [*] No suitable checkpoint!')
 93 |         return False
 94 | 
 95 | def predict():
 96 |     is_training = False
 97 |     
 98 |     with tf.device('/gpu:'+str(gpu_to_use)):
 99 |         pointgrid_ph, cat_label_ph, seg_label_ph = placeholder_inputs()
100 |         is_training_ph = tf.placeholder(tf.bool, shape=())
101 | 
102 |         # model
103 |         pred_cat, pred_seg = model.get_model(pointgrid_ph, is_training=is_training_ph)
104 | 
105 |     # Add ops to save and restore all the variables.
106 |     saver = tf.train.Saver()
107 | 
108 |     # Later, launch the model, use the saver to restore variables from disk, and
109 |     # do some work with the model.
110 |     
111 |     config = tf.ConfigProto()
112 |     config.gpu_options.allow_growth = True
113 |     config.allow_soft_placement = True
114 | 
115 |     with tf.Session(config=config) as sess:
116 |         if not os.path.exists(output_dir):
117 |             os.mkdir(output_dir)
118 | 
119 |         flog = open(os.path.join(output_dir, 'log.txt'), 'w')
120 | 
121 |         # Restore variables from disk.
122 |         ckpt_dir = './train_results/trained_models'
123 |         if not load_checkpoint(ckpt_dir, sess):
124 |             exit()
125 | 
126 |         if not os.path.exists('../data/ShapeNet/test-PointGrid'):
127 |             os.mkdir('../data/ShapeNet/test-PointGrid')
128 |         
129 |         avg_cat_accuracy = 0.0
130 |         cat_accuracy = np.zeros((model.NUM_CATEGORY), dtype=np.float32)
131 |         cat_obj = np.zeros((model.NUM_CATEGORY), dtype=np.int32)
132 |         avg_iou = 0.0
133 |         cat_iou = np.zeros((model.NUM_CATEGORY), dtype=np.float32)
134 |         for loop in range(len(TESTING_FILE_LIST)):
135 |             mat_content = scipy.io.loadmat('../data/ShapeNet/test/' + TESTING_FILE_LIST[loop] + '.mat')
136 |             pc = mat_content['points']
137 |             labels = np.squeeze(mat_content['labels'])
138 |             category = mat_content['category'][0][0]
139 |             cat_label = model.integer_label_to_one_hot_label(category)
140 |             category = int(category)
141 |             cat_obj[category] += 1
142 |             seg_label = model.integer_label_to_one_hot_label(labels)
143 |             pointgrid, pointgrid_label, index = model.pc2voxel(pc, seg_label)
144 |             pointgrid = np.expand_dims(pointgrid, axis=0)
145 |             cat_label = np.expand_dims(cat_label, axis=0)
146 |             pointgrid_label = np.expand_dims(pointgrid_label, axis=0)
147 |             feed_dict = {
148 |                          pointgrid_ph: pointgrid,
149 |                          cat_label_ph: cat_label,
150 |                          seg_label_ph: pointgrid_label,
151 |                          is_training_ph: is_training,
152 |                         }
153 |             pred_cat_val, pred_seg_val = sess.run([pred_cat, pred_seg], feed_dict = feed_dict)
154 |             pred_cat_val = np.argmax(pred_cat_val[0, :], axis=0)
155 |             pred_seg_val = pred_seg_val[0, :, :, :, :, :]
156 |             avg_cat_accuracy += (pred_cat_val == category)
157 |             cat_accuracy[category] += (pred_cat_val == category)
158 |             pred_point_label = model.populateOneHotSegLabel(pc, pred_seg_val, index)
159 |             if purify == True:
160 |                 pre_label = pred_point_label
161 |                 for i in range(pc.shape[0]):
162 |                     idx = np.argsort(np.sum((pc[i, :] - pc) ** 2, axis=1))
163 |                     j, L = 0, []
164 |                     for _ in range(knn):
165 |                         if (idx[j] == i):
166 |                             j += 1
167 |                         L.append(pre_label[idx[j]])
168 |                         j += 1
169 |                     majority = max(set(L), key=L.count)
170 |                     if (pre_label[i] == 0 or len(set(L)) == 1):
171 |                         pred_point_label[i] = majority
172 |             iou = model.intersection_over_union(pred_point_label, labels)
173 |             avg_iou += iou
174 |             cat_iou[category] += iou
175 |             if (cat_obj[category] <= 3):
176 |                 output_color_point_cloud(pc, pred_point_label, '../data/ShapeNet/test-PointGrid/' + category2name[category] + '_' + str(cat_obj[category]) + '.obj')
177 |             printout(flog, '%d/%d %s' % ((loop+1), len(TESTING_FILE_LIST), TESTING_FILE_LIST[loop]))
178 |             printout(flog, '----------')
179 | 
180 |         avg_cat_accuracy /= float(np.sum(cat_obj))
181 |         avg_iou /= float(np.sum(cat_obj))
182 |         printout(flog, 'Average classification accuracy: %f' % avg_cat_accuracy)
183 |         printout(flog, 'Average IoU: %f' % avg_iou)
184 |         printout(flog, 'CategoryName, CategorySize, ClassificationAccuracy, SegmentationIoU')
185 |         for i in range(model.NUM_CATEGORY):
186 |             cat_accuracy[i] /= float(cat_obj[i])
187 |             cat_iou[i] /= float(cat_obj[i])
188 |             printout(flog, '\t%s (%d): %f, %f' % (category2name[i], cat_obj[i], cat_accuracy[i], cat_iou[i]))
189 | 
190 | 
191 | 
192 | with tf.Graph().as_default():
193 |     predict()
194 | 


--------------------------------------------------------------------------------
/code/train.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import subprocess
  3 | import tensorflow as tf
  4 | import threading
  5 | import numpy as np
  6 | import scipy.io
  7 | from datetime import datetime
  8 | import json
  9 | import os
 10 | import sys
 11 | import glob
 12 | BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 13 | sys.path.append(BASE_DIR)
 14 | sys.path.append(os.path.dirname(BASE_DIR))
 15 | import network as model
 16 | 
 17 | # DEFAULT SETTINGS
 18 | parser = argparse.ArgumentParser()
 19 | parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]')
 20 | parser.add_argument('--batch', type=int, default=32, help='Batch Size during training [default: 32]')
 21 | parser.add_argument('--epoch', type=int, default=200, help='Epoch to run [default: 200]')
 22 | parser.add_argument('--output_dir', type=str, default='train_results', help='Directory that stores all training logs and trained models')
 23 | parser.add_argument('--wd', type=float, default=0, help='Weight Decay [Default: 0.0]')
 24 | FLAGS = parser.parse_args()
 25 | 
 26 | # MAIN SCRIPT
 27 | batch_size = FLAGS.batch
 28 | output_dir = FLAGS.output_dir
 29 | 
 30 | if not os.path.exists(output_dir):
 31 |     os.mkdir(output_dir)
 32 | 
 33 | print('#### Batch Size: {0}'.format(batch_size))
 34 | print('#### Training using GPU: {0}'.format(FLAGS.gpu))
 35 | 
 36 | LEARNING_RATE = 1e-4
 37 | TRAINING_EPOCHES = FLAGS.epoch
 38 | print('### Training epoch: {0}'.format(TRAINING_EPOCHES))
 39 | 
 40 | def get_file_name(file_path):
 41 |     parts = file_path.split('/')
 42 |     part = parts[-1]
 43 |     parts = part.split('.')
 44 |     return parts[0]
 45 | 
 46 | TRAINING_FILE_LIST = [get_file_name(file_name) for file_name in glob.glob('../data/ShapeNet/train/' + '*.mat')]
 47 | 
 48 | MODEL_STORAGE_PATH = os.path.join(output_dir, 'trained_models')
 49 | if not os.path.exists(MODEL_STORAGE_PATH):
 50 |     os.mkdir(MODEL_STORAGE_PATH)
 51 | 
 52 | LOG_STORAGE_PATH = os.path.join(output_dir, 'logs')
 53 | if not os.path.exists(LOG_STORAGE_PATH):
 54 |     os.mkdir(LOG_STORAGE_PATH)
 55 | 
 56 | SUMMARIES_FOLDER =  os.path.join(output_dir, 'summaries')
 57 | if not os.path.exists(SUMMARIES_FOLDER):
 58 |     os.mkdir(SUMMARIES_FOLDER)
 59 | 
 60 | def printout(flog, data):
 61 |     print(data)
 62 |     flog.write(data + '\n')
 63 | 
 64 | def load_and_enqueue(sess, enqueue_op, pointgrid_ph, cat_label_ph, seg_label_ph):
 65 |     for epoch in range(1000 * TRAINING_EPOCHES):
 66 |         train_file_idx = np.arange(0, len(TRAINING_FILE_LIST))
 67 |         np.random.shuffle(train_file_idx)
 68 |         for loop in range(len(TRAINING_FILE_LIST)):
 69 |             mat_content = scipy.io.loadmat('../data/ShapeNet/train/' + TRAINING_FILE_LIST[train_file_idx[loop]] + '.mat')
 70 |             pc = mat_content['points']
 71 |             labels = np.squeeze(mat_content['labels'])
 72 |             category = mat_content['category'][0][0]
 73 |             pc = model.rotate_pc(pc)
 74 |             cat_label = model.integer_label_to_one_hot_label(category)
 75 |             seg_label = model.integer_label_to_one_hot_label(labels)
 76 |             pointgrid, pointgrid_label, _ = model.pc2voxel(pc, seg_label)
 77 |             sess.run(enqueue_op, feed_dict={pointgrid_ph: pointgrid, cat_label_ph: cat_label, seg_label_ph: pointgrid_label})
 78 | 
 79 | def placeholder_inputs():
 80 |     pointgrid_ph = tf.placeholder(tf.float32, shape=(model.N, model.N, model.N, model.NUM_FEATURES))
 81 |     cat_label_ph = tf.placeholder(tf.float32, shape=(model.NUM_CATEGORY))
 82 |     seg_label_ph = tf.placeholder(tf.float32, shape=(model.N, model.N, model.N, model.K+1, model.NUM_SEG_PART))
 83 |     return pointgrid_ph, cat_label_ph, seg_label_ph
 84 | 
 85 | def load_checkpoint(checkpoint_dir, session, var_list=None):
 86 |     print(' [*] Loading checkpoint...')
 87 |     ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
 88 |     if ckpt and ckpt.model_checkpoint_path:
 89 |         ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
 90 |         ckpt_path = os.path.join(checkpoint_dir, ckpt_name)
 91 |     try:
 92 |         restorer = tf.train.Saver(var_list)
 93 |         restorer.restore(session, ckpt_path)
 94 |         print(' [*] Loading successful! Copy variables from % s' % ckpt_path)
 95 |         return True
 96 |     except:
 97 |         print(' [*] No suitable checkpoint!')
 98 |         return False
 99 | 
100 | class StoppableThread(threading.Thread):
101 |     """Thread class with a stop() method. The thread itself has to check
102 |     regularly for the stopped() condition."""
103 | 
104 |     def __init__(self, target=None, args=None):
105 |         super(StoppableThread, self).__init__(target=target, args=args)
106 |         self._stop_event = threading.Event()
107 | 
108 |     def stop(self):
109 |         self._stop_event.set()
110 | 
111 |     def stopped(self):
112 |         return self._stop_event.is_set()
113 | 
114 | def train():
115 |     with tf.Graph().as_default():
116 |         with tf.device('/gpu:'+str(FLAGS.gpu)):
117 |             pointgrid_ph, cat_label_ph, seg_label_ph = placeholder_inputs()
118 |             is_training_ph = tf.placeholder(tf.bool, shape=())
119 | 
120 |             queue = tf.FIFOQueue(capacity=20*batch_size, dtypes=[tf.float32, tf.float32, tf.float32],\
121 |                                                          shapes=[[model.N, model.N, model.N, model.NUM_FEATURES],\
122 |                                                                  [model.NUM_CATEGORY],
123 |                                                                  [model.N, model.N, model.N, model.K+1, model.NUM_SEG_PART]])
124 |             enqueue_op = queue.enqueue([pointgrid_ph, cat_label_ph, seg_label_ph])
125 |             dequeue_pointgrid, dequeue_cat_label, dequeue_seg_label = queue.dequeue_many(batch_size)
126 | 
127 |             # model
128 |             pred_cat, pred_seg = model.get_model(dequeue_pointgrid, is_training=is_training_ph)
129 | 
130 |             # loss
131 |             total_loss, cat_loss, seg_loss = model.get_loss(pred_cat, dequeue_cat_label, pred_seg, dequeue_seg_label)
132 | 
133 |             # optimization
134 |             total_var = tf.trainable_variables()
135 |             step = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE).minimize(total_loss, var_list=total_var)
136 | 
137 |         # write logs to the disk
138 |         flog = open(os.path.join(LOG_STORAGE_PATH, 'log.txt'), 'w')
139 | 
140 |         saver = tf.train.Saver()
141 | 
142 |         config = tf.ConfigProto()
143 |         config.gpu_options.allow_growth = True
144 |         config.allow_soft_placement = True
145 |         sess = tf.Session(config=config)
146 | 
147 |         ckpt_dir = './train_results/trained_models'
148 |         if not load_checkpoint(ckpt_dir, sess):
149 |             sess.run(tf.global_variables_initializer())
150 | 
151 |         train_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/train', sess.graph)
152 |         test_writer = tf.summary.FileWriter(SUMMARIES_FOLDER + '/test')
153 | 
154 |         fcmd = open(os.path.join(LOG_STORAGE_PATH, 'cmd.txt'), 'w')
155 |         fcmd.write(str(FLAGS))
156 |         fcmd.close()
157 | 
158 |         def train_one_epoch(epoch_num):
159 |             is_training = True
160 | 
161 |             num_data = len(TRAINING_FILE_LIST)
162 |             num_batch = num_data // batch_size
163 |             total_loss_acc = 0.0
164 |             cat_loss_acc = 0.0
165 |             seg_loss_acc = 0.0
166 |             display_mark = max([num_batch // 4, 1])
167 |             for i in range(num_batch):
168 |                 _, total_loss_val, cat_loss_val, seg_loss_val = sess.run([step, total_loss, cat_loss, seg_loss], feed_dict={is_training_ph: is_training})
169 |                 total_loss_acc += total_loss_val
170 |                 cat_loss_acc += cat_loss_val
171 |                 seg_loss_acc += seg_loss_val
172 | 
173 |                 if ((i+1) % display_mark == 0):
174 |                     printout(flog, 'Epoch %d/%d - Iter %d/%d' % (epoch_num+1, TRAINING_EPOCHES, i+1, num_batch))
175 |                     printout(flog, 'Total Loss: %f' % (total_loss_acc / (i+1)))
176 |                     printout(flog, 'Classification Loss: %f' % (cat_loss_acc / (i+1)))
177 |                     printout(flog, 'Segmentation Loss: %f' % (seg_loss_acc / (i+1)))
178 | 
179 |             printout(flog, '\tMean Total Loss: %f' % (total_loss_acc / num_batch))
180 |             printout(flog, '\tMean Classification Loss: %f' % (cat_loss_acc / num_batch))
181 |             printout(flog, '\tMean Segmentation Loss: %f' % (seg_loss_acc / num_batch))
182 | 
183 |         if not os.path.exists(MODEL_STORAGE_PATH):
184 |             os.mkdir(MODEL_STORAGE_PATH)
185 | 
186 |         coord = tf.train.Coordinator()
187 |         for num_thread in range(16):
188 |             t = StoppableThread(target=load_and_enqueue, args=(sess, enqueue_op, pointgrid_ph, cat_label_ph, seg_label_ph))
189 |             t.setDaemon(True)
190 |             t.start()
191 |             coord.register_thread(t)
192 | 
193 |         for epoch in range(TRAINING_EPOCHES):
194 |             printout(flog, '\n>>> Training for the epoch %d/%d ...' % (epoch+1, TRAINING_EPOCHES))
195 | 
196 |             train_one_epoch(epoch)
197 | 
198 |             if (epoch+1) % 1 == 0:
199 |                 cp_filename = saver.save(sess, os.path.join(MODEL_STORAGE_PATH, 'epoch_' + str(epoch+1)+'.ckpt'))
200 |                 printout(flog, 'Successfully store the checkpoint model into ' + cp_filename)
201 | 
202 |             flog.flush()
203 |         flog.close()
204 | 
205 | if __name__=='__main__':
206 |     train()
207 | 


--------------------------------------------------------------------------------
/code/train_results/trained_models/readme.txt:
--------------------------------------------------------------------------------
1 | Please download the trained model here.
2 | https://drive.google.com/file/d/1-He7KBxocvrOt2O7RYdFJaanzYhg8Qvs/view?usp=sharing
3 | 


--------------------------------------------------------------------------------
/code/utils/tf_util.py:
--------------------------------------------------------------------------------
  1 | """ Wrapper functions for TensorFlow layers.
  2 | 
  3 | Author: Truc D. Le
  4 | Date: June 2018
  5 | Author: Charles R. Qi
  6 | Date: November 2016
  7 | """
  8 | 
  9 | import numpy as np
 10 | from functools import partial
 11 | import tensorflow as tf
 12 | import tensorflow.contrib.slim as slim
 13 | 
 14 | batch_norm = partial(slim.batch_norm, decay=0.95, scale=True, epsilon=1e-5, updates_collections=None)
 15 | 
 16 | def _variable_on_cpu(name, shape, initializer, use_fp16=False):
 17 |   """Helper to create a Variable stored on CPU memory.
 18 |   Args:
 19 |     name: name of the variable
 20 |     shape: list of ints
 21 |     initializer: initializer for Variable
 22 |   Returns:
 23 |     Variable Tensor
 24 |   """
 25 |   with tf.device('/cpu:0'):
 26 |     dtype = tf.float16 if use_fp16 else tf.float32
 27 |     var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype)
 28 |   return var
 29 | 
 30 | def _variable_with_weight_decay(name, shape, stddev, wd, use_xavier=True):
 31 |   """Helper to create an initialized Variable with weight decay.
 32 | 
 33 |   Note that the Variable is initialized with a truncated normal distribution.
 34 |   A weight decay is added only if one is specified.
 35 | 
 36 |   Args:
 37 |     name: name of the variable
 38 |     shape: list of ints
 39 |     stddev: standard deviation of a truncated Gaussian
 40 |     wd: add L2Loss weight decay multiplied by this float. If None, weight
 41 |         decay is not added for this Variable.
 42 |     use_xavier: bool, whether to use xavier initializer
 43 | 
 44 |   Returns:
 45 |     Variable Tensor
 46 |   """
 47 |   if use_xavier:
 48 |     initializer = tf.contrib.layers.xavier_initializer()
 49 |   else:
 50 |     initializer = tf.truncated_normal_initializer(stddev=stddev)
 51 |   var = _variable_on_cpu(name, shape, initializer)
 52 |   if wd is not None:
 53 |     weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss')
 54 |     tf.add_to_collection('losses', weight_decay)
 55 |   return var
 56 | 
 57 | 
 58 | def conv1d(inputs,
 59 |            num_output_channels,
 60 |            kernel_size,
 61 |            scope,
 62 |            stride=1,
 63 |            padding='SAME',
 64 |            use_xavier=True,
 65 |            stddev=1e-3,
 66 |            weight_decay=0.0,
 67 |            activation_fn=tf.nn.relu,
 68 |            bn=False,
 69 |            bn_decay=None,
 70 |            is_training=None):
 71 |   """ 1D convolution with non-linear operation.
 72 | 
 73 |   Args:
 74 |     inputs: 3-D tensor variable BxLxC
 75 |     num_output_channels: int
 76 |     kernel_size: int
 77 |     scope: string
 78 |     stride: int
 79 |     padding: 'SAME' or 'VALID'
 80 |     use_xavier: bool, use xavier_initializer if true
 81 |     stddev: float, stddev for truncated_normal init
 82 |     weight_decay: float
 83 |     activation_fn: function
 84 |     bn: bool, whether to use batch norm
 85 |     bn_decay: float or float tensor variable in [0,1]
 86 |     is_training: bool Tensor variable
 87 | 
 88 |   Returns:
 89 |     Variable tensor
 90 |   """
 91 |   with tf.variable_scope(scope) as sc:
 92 |     num_in_channels = inputs.get_shape()[-1].value
 93 |     kernel_shape = [kernel_size,
 94 |                     num_in_channels, num_output_channels]
 95 |     kernel = _variable_with_weight_decay('weights',
 96 |                                          shape=kernel_shape,
 97 |                                          use_xavier=use_xavier,
 98 |                                          stddev=stddev,
 99 |                                          wd=weight_decay)
100 |     outputs = tf.nn.conv1d(inputs, kernel,
101 |                            stride=stride,
102 |                            padding=padding)
103 |     biases = _variable_on_cpu('biases', [num_output_channels],
104 |                               tf.constant_initializer(0.0))
105 |     outputs = tf.nn.bias_add(outputs, biases)
106 | 
107 |     if bn:
108 |       outputs = batch_norm(outputs, is_training=is_training)
109 |       #outputs = batch_norm_for_conv1d(outputs, is_training,
110 |       #                                bn_decay=bn_decay, scope='bn')
111 | 
112 |     if activation_fn is not None:
113 |       outputs = activation_fn(outputs)
114 |     return outputs
115 | 
116 | 
117 | 
118 | 
119 | def conv2d(inputs,
120 |            num_output_channels,
121 |            kernel_size,
122 |            scope,
123 |            stride=[1, 1],
124 |            padding='SAME',
125 |            use_xavier=True,
126 |            stddev=1e-3,
127 |            weight_decay=0.0,
128 |            activation_fn=tf.nn.relu,
129 |            bn=False,
130 |            bn_decay=None,
131 |            is_training=None):
132 |   """ 2D convolution with non-linear operation.
133 | 
134 |   Args:
135 |     inputs: 4-D tensor variable BxHxWxC
136 |     num_output_channels: int
137 |     kernel_size: a list of 2 ints
138 |     scope: string
139 |     stride: a list of 2 ints
140 |     padding: 'SAME' or 'VALID'
141 |     use_xavier: bool, use xavier_initializer if true
142 |     stddev: float, stddev for truncated_normal init
143 |     weight_decay: float
144 |     activation_fn: function
145 |     bn: bool, whether to use batch norm
146 |     bn_decay: float or float tensor variable in [0,1]
147 |     is_training: bool Tensor variable
148 | 
149 |   Returns:
150 |     Variable tensor
151 |   """
152 |   with tf.variable_scope(scope) as sc:
153 |       kernel_h, kernel_w = kernel_size
154 |       num_in_channels = inputs.get_shape()[-1].value
155 |       kernel_shape = [kernel_h, kernel_w,
156 |                       num_in_channels, num_output_channels]
157 |       kernel = _variable_with_weight_decay('weights',
158 |                                            shape=kernel_shape,
159 |                                            use_xavier=use_xavier,
160 |                                            stddev=stddev,
161 |                                            wd=weight_decay)
162 |       stride_h, stride_w = stride
163 |       outputs = tf.nn.conv2d(inputs, kernel,
164 |                              [1, stride_h, stride_w, 1],
165 |                              padding=padding)
166 |       biases = _variable_on_cpu('biases', [num_output_channels],
167 |                                 tf.constant_initializer(0.0))
168 |       outputs = tf.nn.bias_add(outputs, biases)
169 | 
170 |       if bn:
171 |         outputs = batch_norm(outputs, is_training=is_training)
172 |         #outputs = batch_norm_for_conv2d(outputs, is_training,
173 |         #                                bn_decay=bn_decay, scope='bn')
174 | 
175 |       if activation_fn is not None:
176 |         outputs = activation_fn(outputs)
177 |       return outputs
178 | 
179 | 
180 | def conv2d_transpose(inputs,
181 |                      num_output_channels,
182 |                      kernel_size,
183 |                      scope,
184 |                      stride=[1, 1],
185 |                      padding='SAME',
186 |                      use_xavier=True,
187 |                      stddev=1e-3,
188 |                      weight_decay=0.0,
189 |                      activation_fn=tf.nn.relu,
190 |                      bn=False,
191 |                      bn_decay=None,
192 |                      is_training=None):
193 |   """ 2D convolution transpose with non-linear operation.
194 | 
195 |   Args:
196 |     inputs: 4-D tensor variable BxHxWxC
197 |     num_output_channels: int
198 |     kernel_size: a list of 2 ints
199 |     scope: string
200 |     stride: a list of 2 ints
201 |     padding: 'SAME' or 'VALID'
202 |     use_xavier: bool, use xavier_initializer if true
203 |     stddev: float, stddev for truncated_normal init
204 |     weight_decay: float
205 |     activation_fn: function
206 |     bn: bool, whether to use batch norm
207 |     bn_decay: float or float tensor variable in [0,1]
208 |     is_training: bool Tensor variable
209 | 
210 |   Returns:
211 |     Variable tensor
212 | 
213 |   Note: conv2d(conv2d_transpose(a, num_out, ksize, stride), a.shape[-1], ksize, stride) == a
214 |   """
215 |   with tf.variable_scope(scope) as sc:
216 |       kernel_h, kernel_w = kernel_size
217 |       num_in_channels = inputs.get_shape()[-1].value
218 |       kernel_shape = [kernel_h, kernel_w,
219 |                       num_output_channels, num_in_channels] # reversed to conv2d
220 |       kernel = _variable_with_weight_decay('weights',
221 |                                            shape=kernel_shape,
222 |                                            use_xavier=use_xavier,
223 |                                            stddev=stddev,
224 |                                            wd=weight_decay)
225 |       stride_h, stride_w = stride
226 |       
227 |       # from slim.convolution2d_transpose
228 |       def get_deconv_dim(dim_size, stride_size, kernel_size, padding):
229 |           dim_size *= stride_size
230 | 
231 |           if padding == 'VALID' and dim_size is not None:
232 |             dim_size += max(kernel_size - stride_size, 0)
233 |           return dim_size
234 | 
235 |       # caculate output shape
236 |       batch_size = inputs.get_shape()[0].value
237 |       height = inputs.get_shape()[1].value
238 |       width = inputs.get_shape()[2].value
239 |       out_height = get_deconv_dim(height, stride_h, kernel_h, padding)
240 |       out_width = get_deconv_dim(width, stride_w, kernel_w, padding)
241 |       output_shape = [batch_size, out_height, out_width, num_output_channels]
242 | 
243 |       outputs = tf.nn.conv2d_transpose(inputs, kernel, output_shape,
244 |                              [1, stride_h, stride_w, 1],
245 |                              padding=padding)
246 |       biases = _variable_on_cpu('biases', [num_output_channels],
247 |                                 tf.constant_initializer(0.0))
248 |       outputs = tf.nn.bias_add(outputs, biases)
249 | 
250 |       if bn:
251 |         outputs = batch_norm(outputs, is_training=is_training)
252 |         #outputs = batch_norm_for_conv2d(outputs, is_training,
253 |         #                                bn_decay=bn_decay, scope='bn')
254 | 
255 |       if activation_fn is not None:
256 |         outputs = activation_fn(outputs)
257 |       return outputs
258 | 
259 |    
260 | 
261 | def conv3d(inputs,
262 |            num_output_channels,
263 |            kernel_size,
264 |            scope,
265 |            stride=[1, 1, 1],
266 |            padding='SAME',
267 |            use_xavier=True,
268 |            stddev=1e-3,
269 |            weight_decay=0.0,
270 |            activation_fn=tf.nn.relu,
271 |            bn=False,
272 |            bn_decay=None,
273 |            is_training=None):
274 |   """ 3D convolution with non-linear operation.
275 | 
276 |   Args:
277 |     inputs: 5-D tensor variable BxDxHxWxC
278 |     num_output_channels: int
279 |     kernel_size: a list of 3 ints
280 |     scope: string
281 |     stride: a list of 3 ints
282 |     padding: 'SAME' or 'VALID'
283 |     use_xavier: bool, use xavier_initializer if true
284 |     stddev: float, stddev for truncated_normal init
285 |     weight_decay: float
286 |     activation_fn: function
287 |     bn: bool, whether to use batch norm
288 |     bn_decay: float or float tensor variable in [0,1]
289 |     is_training: bool Tensor variable
290 | 
291 |   Returns:
292 |     Variable tensor
293 |   """
294 |   with tf.variable_scope(scope) as sc:
295 |     kernel_d, kernel_h, kernel_w = kernel_size
296 |     num_in_channels = inputs.get_shape()[-1].value
297 |     kernel_shape = [kernel_d, kernel_h, kernel_w,
298 |                     num_in_channels, num_output_channels]
299 |     kernel = _variable_with_weight_decay('weights',
300 |                                          shape=kernel_shape,
301 |                                          use_xavier=use_xavier,
302 |                                          stddev=stddev,
303 |                                          wd=weight_decay)
304 |     stride_d, stride_h, stride_w = stride
305 |     outputs = tf.nn.conv3d(inputs, kernel,
306 |                            [1, stride_d, stride_h, stride_w, 1],
307 |                            padding=padding)
308 |     biases = _variable_on_cpu('biases', [num_output_channels],
309 |                               tf.constant_initializer(0.0))
310 |     outputs = tf.nn.bias_add(outputs, biases)
311 |     
312 |     if bn:
313 |       outputs = batch_norm(outputs, is_training=is_training)
314 |       #outputs = batch_norm_for_conv3d(outputs, is_training,
315 |       #                                bn_decay=bn_decay, scope='bn')
316 | 
317 |     if activation_fn is not None:
318 |       outputs = activation_fn(outputs)
319 |     return outputs
320 | 
321 | 
322 | def conv3d_transpose(inputs,
323 |                      num_output_channels,
324 |                      kernel_size,
325 |                      scope,
326 |                      stride=[1, 1, 1],
327 |                      padding='SAME',
328 |                      use_xavier=True,
329 |                      stddev=1e-3,
330 |                      weight_decay=0.0,
331 |                      activation_fn=tf.nn.relu,
332 |                      bn=False,
333 |                      bn_decay=None,
334 |                      is_training=None):
335 |   """ 3D convolution transpose with non-linear operation.
336 | 
337 |   Args:
338 |     inputs: 5-D tensor variable BxDxHxWxC
339 |     num_output_channels: int
340 |     kernel_size: a list of 3 ints
341 |     scope: string
342 |     stride: a list of 3 ints
343 |     padding: 'SAME' or 'VALID'
344 |     use_xavier: bool, use xavier_initializer if true
345 |     stddev: float, stddev for truncated_normal init
346 |     weight_decay: float
347 |     activation_fn: function
348 |     bn: bool, whether to use batch norm
349 |     bn_decay: float or float tensor variable in [0,1]
350 |     is_training: bool Tensor variable
351 | 
352 |   Returns:
353 |     Variable tensor
354 | 
355 |   Note: conv3d(conv3d_transpose(a, num_out, ksize, stride), a.shape[-1], ksize, stride) == a
356 |   """
357 |   with tf.variable_scope(scope) as sc:
358 |       kernel_d, kernel_h, kernel_w = kernel_size
359 |       num_in_channels = inputs.get_shape()[-1].value
360 |       kernel_shape = [kernel_d, kernel_h, kernel_w,
361 |                       num_output_channels, num_in_channels] # reversed to conv3d
362 |       kernel = _variable_with_weight_decay('weights',
363 |                                            shape=kernel_shape,
364 |                                            use_xavier=use_xavier,
365 |                                            stddev=stddev,
366 |                                            wd=weight_decay)
367 |       stride_d, stride_h, stride_w = stride
368 |       
369 |       # from slim.convolution3d_transpose
370 |       def get_deconv_dim(dim_size, stride_size, kernel_size, padding):
371 |           dim_size *= stride_size
372 | 
373 |           if padding == 'VALID' and dim_size is not None:
374 |             dim_size += max(kernel_size - stride_size, 0)
375 |           return dim_size
376 | 
377 |       # caculate output shape
378 |       batch_size = inputs.get_shape()[0].value
379 |       depth = inputs.get_shape()[1].value
380 |       height = inputs.get_shape()[2].value
381 |       width = inputs.get_shape()[3].value
382 |       out_depth = get_deconv_dim(depth, stride_d, kernel_d, padding)
383 |       out_height = get_deconv_dim(height, stride_h, kernel_h, padding)
384 |       out_width = get_deconv_dim(width, stride_w, kernel_w, padding)
385 |       output_shape = [batch_size, out_depth, out_height, out_width, num_output_channels]
386 | 
387 |       outputs = tf.nn.conv3d_transpose(inputs, kernel, output_shape,
388 |                              [1, stride_d, stride_h, stride_w, 1],
389 |                              padding=padding)
390 |       biases = _variable_on_cpu('biases', [num_output_channels],
391 |                                 tf.constant_initializer(0.0))
392 |       outputs = tf.nn.bias_add(outputs, biases)
393 | 
394 |       if bn:
395 |         outputs = batch_norm(outputs, is_training=is_training)
396 |         #outputs = batch_norm_for_conv3d(outputs, is_training,
397 |         #                                bn_decay=bn_decay, scope='bn')
398 | 
399 |       if activation_fn is not None:
400 |         outputs = activation_fn(outputs)
401 |       return outputs
402 | 
403 | 
404 | 
405 | def fully_connected(inputs,
406 |                     num_outputs,
407 |                     scope,
408 |                     use_xavier=True,
409 |                     stddev=1e-3,
410 |                     weight_decay=0.0,
411 |                     activation_fn=tf.nn.relu,
412 |                     bn=False,
413 |                     bn_decay=None,
414 |                     is_training=None):
415 |   """ Fully connected layer with non-linear operation.
416 |   
417 |   Args:
418 |     inputs: 2-D tensor BxN
419 |     num_outputs: int
420 |   
421 |   Returns:
422 |     Variable tensor of size B x num_outputs.
423 |   """
424 |   with tf.variable_scope(scope) as sc:
425 |     num_input_units = inputs.get_shape()[-1].value
426 |     weights = _variable_with_weight_decay('weights',
427 |                                           shape=[num_input_units, num_outputs],
428 |                                           use_xavier=use_xavier,
429 |                                           stddev=stddev,
430 |                                           wd=weight_decay)
431 |     outputs = tf.matmul(inputs, weights)
432 |     biases = _variable_on_cpu('biases', [num_outputs],
433 |                              tf.constant_initializer(0.0))
434 |     outputs = tf.nn.bias_add(outputs, biases)
435 |      
436 |     if bn:
437 |       outputs = batch_norm(outputs, is_training=is_training)
438 |       #outputs = batch_norm_for_fc(outputs, is_training, bn_decay, 'bn')
439 | 
440 |     if activation_fn is not None:
441 |       outputs = activation_fn(outputs)
442 |     return outputs
443 | 
444 | 
445 | def max_pool2d(inputs,
446 |                kernel_size,
447 |                scope,
448 |                stride=[2, 2],
449 |                padding='VALID'):
450 |   """ 2D max pooling.
451 | 
452 |   Args:
453 |     inputs: 4-D tensor BxHxWxC
454 |     kernel_size: a list of 2 ints
455 |     stride: a list of 2 ints
456 |   
457 |   Returns:
458 |     Variable tensor
459 |   """
460 |   with tf.variable_scope(scope) as sc:
461 |     kernel_h, kernel_w = kernel_size
462 |     stride_h, stride_w = stride
463 |     outputs = tf.nn.max_pool(inputs,
464 |                              ksize=[1, kernel_h, kernel_w, 1],
465 |                              strides=[1, stride_h, stride_w, 1],
466 |                              padding=padding,
467 |                              name=sc.name)
468 |     return outputs
469 | 
470 | def avg_pool2d(inputs,
471 |                kernel_size,
472 |                scope,
473 |                stride=[2, 2],
474 |                padding='VALID'):
475 |   """ 2D avg pooling.
476 | 
477 |   Args:
478 |     inputs: 4-D tensor BxHxWxC
479 |     kernel_size: a list of 2 ints
480 |     stride: a list of 2 ints
481 |   
482 |   Returns:
483 |     Variable tensor
484 |   """
485 |   with tf.variable_scope(scope) as sc:
486 |     kernel_h, kernel_w = kernel_size
487 |     stride_h, stride_w = stride
488 |     outputs = tf.nn.avg_pool(inputs,
489 |                              ksize=[1, kernel_h, kernel_w, 1],
490 |                              strides=[1, stride_h, stride_w, 1],
491 |                              padding=padding,
492 |                              name=sc.name)
493 |     return outputs
494 | 
495 | 
496 | def max_pool3d(inputs,
497 |                kernel_size,
498 |                scope,
499 |                stride=[2, 2, 2],
500 |                padding='VALID'):
501 |   """ 3D max pooling.
502 | 
503 |   Args:
504 |     inputs: 5-D tensor BxDxHxWxC
505 |     kernel_size: a list of 3 ints
506 |     stride: a list of 3 ints
507 |   
508 |   Returns:
509 |     Variable tensor
510 |   """
511 |   with tf.variable_scope(scope) as sc:
512 |     kernel_d, kernel_h, kernel_w = kernel_size
513 |     stride_d, stride_h, stride_w = stride
514 |     outputs = tf.nn.max_pool3d(inputs,
515 |                                ksize=[1, kernel_d, kernel_h, kernel_w, 1],
516 |                                strides=[1, stride_d, stride_h, stride_w, 1],
517 |                                padding=padding,
518 |                                name=sc.name)
519 |     return outputs
520 | 
521 | def avg_pool3d(inputs,
522 |                kernel_size,
523 |                scope,
524 |                stride=[2, 2, 2],
525 |                padding='VALID'):
526 |   """ 3D avg pooling.
527 | 
528 |   Args:
529 |     inputs: 5-D tensor BxDxHxWxC
530 |     kernel_size: a list of 3 ints
531 |     stride: a list of 3 ints
532 |   
533 |   Returns:
534 |     Variable tensor
535 |   """
536 |   with tf.variable_scope(scope) as sc:
537 |     kernel_d, kernel_h, kernel_w = kernel_size
538 |     stride_d, stride_h, stride_w = stride
539 |     outputs = tf.nn.avg_pool3d(inputs,
540 |                                ksize=[1, kernel_d, kernel_h, kernel_w, 1],
541 |                                strides=[1, stride_d, stride_h, stride_w, 1],
542 |                                padding=padding,
543 |                                name=sc.name)
544 |     return outputs
545 | 
546 | 
547 | 
548 | 
549 | 
550 | def batch_norm_template(inputs, is_training, scope, moments_dims, bn_decay):
551 |   """ Batch normalization on convolutional maps and beyond...
552 |   Ref.: http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow
553 |   
554 |   Args:
555 |       inputs:        Tensor, k-D input ... x C could be BC or BHWC or BDHWC
556 |       is_training:   boolean tf.Varialbe, true indicates training phase
557 |       scope:         string, variable scope
558 |       moments_dims:  a list of ints, indicating dimensions for moments calculation
559 |       bn_decay:      float or float tensor variable, controling moving average weight
560 |   Return:
561 |       normed:        batch-normalized maps
562 |   """
563 |   with tf.variable_scope(scope) as sc:
564 |     num_channels = inputs.get_shape()[-1].value
565 |     beta = tf.get_variable('beta', shape=[num_channels], initializer=tf.zeros_initializer, dtype=tf.float32, trainable=True)
566 |     gamma = tf.get_variable('gamma', shape=[num_channels], initializer=tf.ones_initializer, dtype=tf.float32, trainable=True)
567 |     batch_mean, batch_var = tf.nn.moments(inputs, moments_dims, name='moments')
568 |     decay = bn_decay if bn_decay is not None else 0.9
569 |     ema = tf.train.ExponentialMovingAverage(decay=decay)
570 |     # Operator that maintains moving averages of variables.
571 |     ema_apply_op = tf.cond(is_training,
572 |                            lambda: ema.apply([batch_mean, batch_var]),
573 |                            lambda: tf.no_op())
574 |     
575 |     # Update moving average and return current batch's avg and var.
576 |     def mean_var_with_update():
577 |       with tf.control_dependencies([ema_apply_op]):
578 |         return tf.identity(batch_mean), tf.identity(batch_var)
579 |     
580 |     # ema.average returns the Variable holding the average of var.
581 |     mean, var = tf.cond(is_training,
582 |                         mean_var_with_update,
583 |                         lambda: (ema.average(batch_mean), ema.average(batch_var)))
584 |     normed = tf.nn.batch_normalization(inputs, mean, var, beta, gamma, 1e-3)
585 |   return normed
586 | 
587 | 
588 | def batch_norm_for_fc(inputs, is_training, bn_decay, scope):
589 |   """ Batch normalization on FC data.
590 |   
591 |   Args:
592 |       inputs:      Tensor, 2D BxC input
593 |       is_training: boolean tf.Varialbe, true indicates training phase
594 |       bn_decay:    float or float tensor variable, controling moving average weight
595 |       scope:       string, variable scope
596 |   Return:
597 |       normed:      batch-normalized maps
598 |   """
599 |   return batch_norm_template(inputs, is_training, scope, [0,], bn_decay)
600 | 
601 | 
602 | def batch_norm_for_conv1d(inputs, is_training, bn_decay, scope):
603 |   """ Batch normalization on 1D convolutional maps.
604 |   
605 |   Args:
606 |       inputs:      Tensor, 3D BLC input maps
607 |       is_training: boolean tf.Varialbe, true indicates training phase
608 |       bn_decay:    float or float tensor variable, controling moving average weight
609 |       scope:       string, variable scope
610 |   Return:
611 |       normed:      batch-normalized maps
612 |   """
613 |   return batch_norm_template(inputs, is_training, scope, [0,1], bn_decay)
614 | 
615 | 
616 | 
617 |   
618 | def batch_norm_for_conv2d(inputs, is_training, bn_decay, scope):
619 |   """ Batch normalization on 2D convolutional maps.
620 |   
621 |   Args:
622 |       inputs:      Tensor, 4D BHWC input maps
623 |       is_training: boolean tf.Varialbe, true indicates training phase
624 |       bn_decay:    float or float tensor variable, controling moving average weight
625 |       scope:       string, variable scope
626 |   Return:
627 |       normed:      batch-normalized maps
628 |   """
629 |   return batch_norm_template(inputs, is_training, scope, [0,1,2], bn_decay)
630 | 
631 | 
632 | 
633 | def batch_norm_for_conv3d(inputs, is_training, bn_decay, scope):
634 |   """ Batch normalization on 3D convolutional maps.
635 |   
636 |   Args:
637 |       inputs:      Tensor, 5D BDHWC input maps
638 |       is_training: boolean tf.Varialbe, true indicates training phase
639 |       bn_decay:    float or float tensor variable, controling moving average weight
640 |       scope:       string, variable scope
641 |   Return:
642 |       normed:      batch-normalized maps
643 |   """
644 |   return batch_norm_template(inputs, is_training, scope, [0,1,2,3], bn_decay)
645 | 
646 | 
647 | def dropout(inputs,
648 |             is_training,
649 |             scope,
650 |             keep_prob=0.5,
651 |             noise_shape=None):
652 |   """ Dropout layer.
653 | 
654 |   Args:
655 |     inputs: tensor
656 |     is_training: boolean tf.Variable
657 |     scope: string
658 |     keep_prob: float in [0,1]
659 |     noise_shape: list of ints
660 | 
661 |   Returns:
662 |     tensor variable
663 |   """
664 |   with tf.variable_scope(scope) as sc:
665 |     outputs = tf.cond(is_training,
666 |                       lambda: tf.nn.dropout(inputs, keep_prob, noise_shape),
667 |                       lambda: inputs)
668 |     return outputs
669 | 


--------------------------------------------------------------------------------
/data/ShapeNet/test/put_test_data_here.txt:
--------------------------------------------------------------------------------
1 | 
2 | 


--------------------------------------------------------------------------------
/data/ShapeNet/train/put_train_data_here.txt:
--------------------------------------------------------------------------------
1 | 
2 | 


--------------------------------------------------------------------------------
/data/put_your_data_here.txt:
--------------------------------------------------------------------------------
1 | 
2 | 


--------------------------------------------------------------------------------