Skip to content

Commit c2b32a8

Browse files
art-programmerart-programmer
committed
first commit
0 parents  commit c2b32a8

31 files changed

+16752
-0
lines changed

‎PlaneSetGeneration/RecordReader.py

+234
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
import tensorflow as tf
2+
import numpy as np
3+
import threading
4+
import PIL.Image as Image
5+
from functools import partial
6+
from multiprocessing import Pool
7+
import cv2
8+
9+
import sys
10+
import os
11+
12+
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
13+
from modules import *
14+
15+
16+
FETCH_BATCH_SIZE=32
17+
BATCH_SIZE=32
18+
HEIGHT=192
19+
WIDTH=256
20+
NUM_PLANES = 50
21+
NUM_THREADS = 4
22+
23+
24+
class RecordReader():
25+
def __init__(self):
26+
return
27+
28+
29+
def getBatch(self, filename_queue, numOutputPlanes = 20, batchSize = BATCH_SIZE, min_after_dequeue = 1000, random=True, getLocal=False, getSegmentation=False, test=True, suffix='forward'):
30+
reader = tf.TFRecordReader()
31+
_, serialized_example = reader.read(filename_queue)
32+
33+
features = tf.parse_single_example(
34+
serialized_example,
35+
# Defaults are not specified since both keys are required.
36+
features={
37+
#'height': tf.FixedLenFeature([], tf.int64),
38+
#'width': tf.FixedLenFeature([], tf.int64),
39+
'image_raw': tf.FixedLenFeature([], tf.string),
40+
'num_planes': tf.FixedLenFeature([], tf.int64),
41+
'plane': tf.FixedLenFeature([NUM_PLANES * 3], tf.float32),
42+
'plane_mask': tf.FixedLenFeature([HEIGHT * WIDTH], tf.int64),
43+
#'validating': tf.FixedLenFeature([], tf.int64)
44+
'depth': tf.FixedLenFeature([HEIGHT * WIDTH], tf.float32),
45+
'normal': tf.FixedLenFeature([HEIGHT * WIDTH * 3], tf.float32),
46+
'boundary_raw': tf.FixedLenFeature([], tf.string),
47+
'grid_s': tf.FixedLenFeature([HEIGHT / 8 * WIDTH / 8 * 1], tf.float32),
48+
'grid_p': tf.FixedLenFeature([HEIGHT / 8 * WIDTH / 8 * 3], tf.float32),
49+
'grid_m_raw': tf.FixedLenFeature([], tf.string),
50+
})
51+
52+
# Convert from a scalar string tensor (whose single string has
53+
# length mnist.IMAGE_PIXELS) to a uint8 tensor with shape
54+
# [mnist.IMAGE_PIXELS].
55+
image = tf.decode_raw(features['image_raw'], tf.uint8)
56+
image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
57+
image = tf.reshape(image, [HEIGHT, WIDTH, 3])
58+
59+
depth = features['depth']
60+
depth = tf.reshape(depth, [HEIGHT, WIDTH, 1])
61+
62+
normal = features['normal']
63+
normal = tf.reshape(normal, [HEIGHT, WIDTH, 3])
64+
65+
66+
# planeAreaThreshold = 12 * 16
67+
# inds, _, counts = tf.unique_with_counts(plane_masks)
68+
# counts = counts * tf.cast(tf.greater(inds, 0), tf.int32)
69+
# numPlanes = tf.minimum(tf.reduce_sum(tf.cast(counts > planeAreaThreshold, tf.int32)), numOutputPlanes)
70+
71+
72+
if '_32' not in suffix and False:
73+
numPlanes = tf.minimum(tf.cast(features['num_planes'], tf.int32), numOutputPlanes)
74+
planes = features['plane']
75+
planes = tf.reshape(planes, [NUM_PLANES, 3])
76+
planes = tf.slice(planes, [0, 0], [numPlanes, 3])
77+
78+
plane_masks = tf.cast(features['plane_mask'], tf.int64)
79+
plane_masks = tf.reshape(plane_masks, [HEIGHT, WIDTH, 1])
80+
plane_mask_array = tf.tile(plane_masks, [1, 1, NUM_PLANES])
81+
coef = tf.range(tf.cast(NUM_PLANES, tf.int64), dtype=tf.int64)
82+
coef = tf.pow(tf.constant(2, tf.int64), coef)
83+
planeMasks = tf.reshape(tf.cast(tf.div(plane_mask_array, coef) % 2, tf.float32), [HEIGHT, WIDTH, NUM_PLANES])
84+
85+
#planeMasks = tf.zeros([HEIGHT, WIDTH, numOutputPlanes])
86+
else:
87+
numPlanes = 30
88+
planes = features['plane']
89+
planes = tf.reshape(planes, [NUM_PLANES, 3])
90+
planes = tf.slice(planes, [0, 0], [numPlanes, 3])
91+
92+
plane_masks = tf.cast(features['plane_mask'], tf.int64)
93+
94+
plane_masks = tf.reshape(plane_masks, [HEIGHT, WIDTH, 1])
95+
plane_mask_array = tf.tile(plane_masks, [1, 1, numPlanes])
96+
coef = tf.range(numPlanes, dtype=tf.int64)
97+
coef = tf.pow(tf.constant(2, tf.int64), coef)
98+
#coef = tf.reshape(tf.matmul(tf.reshape(coef, [-1, numPlanes]), tf.cast(shuffle_inds, tf.int32)), [numPlanes])
99+
#coef = tf.cast(coef, tf.int64)
100+
planeMasks = tf.cast(tf.div(plane_mask_array, coef) % 2, tf.float32)
101+
102+
urange = tf.reshape(tf.range(WIDTH, dtype=tf.float32), [-1, 1])
103+
planeXs = tf.reduce_max(planeMasks, axis=0)
104+
planeMinX = WIDTH - tf.reduce_max(planeXs * (float(WIDTH) - urange), axis=0)
105+
planeMaxX = tf.reduce_max(planeXs * urange, axis=0)
106+
107+
vrange = tf.reshape(tf.range(HEIGHT, dtype=tf.float32), [-1, 1])
108+
planeYs = tf.reduce_max(planeMasks, axis=1)
109+
planeMinY = HEIGHT - tf.reduce_max(planeYs * (float(HEIGHT) - vrange), axis=0)
110+
planeMaxY = tf.reduce_max(planeYs * vrange, axis=0)
111+
112+
planeMaxX = tf.maximum(planeMinX, planeMaxX)
113+
planeMaxY = tf.maximum(planeMinY, planeMaxY)
114+
115+
planeAreas = tf.reduce_sum(planeMasks, axis=[0, 1])
116+
117+
localPlaneWidthThreshold = 32
118+
localPlaneHeightThreshold = 32
119+
globalPlaneAreaThreshold = 16 * 16
120+
globalPlaneWidthThreshold = 8
121+
122+
123+
globalPlaneMask = tf.logical_or(tf.greater(planeMaxX - planeMinX, localPlaneWidthThreshold), tf.greater(planeMaxY - planeMinY, localPlaneHeightThreshold))
124+
globalPlaneMask = tf.logical_and(globalPlaneMask, tf.greater((planeMaxX - planeMinX) * (planeMaxY - planeMinY), globalPlaneAreaThreshold))
125+
globalPlaneMask = tf.logical_and(globalPlaneMask, tf.greater(planeAreas / tf.sqrt(tf.pow(planeMaxX + 1 - planeMinX, 2) + tf.pow(planeMaxY + 1 - planeMinY, 2)), globalPlaneWidthThreshold))
126+
#globalPlaneMask = tf.logical_or(globalPlaneMask, tf.less(tf.range(numPlanes), tf.cast(features['num_planes'], tf.int32)))
127+
#globalPlaneMask = tf.cast(tf.squeeze(globalPlaneMask, axis=[2]), tf.float32)
128+
globalPlaneMask = tf.cast(globalPlaneMask, tf.float32)
129+
130+
weightedPlaneAreas = globalPlaneMask * (planeAreas + HEIGHT * WIDTH) + (1 - globalPlaneMask) * planeAreas
131+
132+
#test = tf.reshape(tf.stack([globalPlaneMask, planeAreas, weightedPlaneAreas, planeMinX, planeMaxX, planeMinY, planeMaxY], axis=0), [7, numPlanes])
133+
134+
planeAreas, sortInds = tf.nn.top_k(weightedPlaneAreas, k=numPlanes)
135+
sortMap = tf.one_hot(sortInds, depth=numPlanes, axis=0)
136+
137+
planeMasks = tf.reshape(tf.matmul(tf.reshape(planeMasks, [HEIGHT * WIDTH, numPlanes]), sortMap), [HEIGHT, WIDTH, numPlanes])
138+
planes = tf.transpose(tf.matmul(planes, sortMap, transpose_a=True), [1, 0])
139+
140+
numPlanes = tf.minimum(tf.cast(tf.round(tf.reduce_sum(globalPlaneMask)), tf.int32), numOutputPlanes)
141+
142+
planes = tf.slice(planes, [0, 0], [numPlanes, 3])
143+
planeMasks = tf.slice(planeMasks, [0, 0, 0], [HEIGHT, WIDTH, numPlanes])
144+
planeMasks = tf.reshape(tf.concat([planeMasks, tf.zeros([HEIGHT, WIDTH, numOutputPlanes - numPlanes])], axis=2), [HEIGHT, WIDTH, numOutputPlanes])
145+
pass
146+
147+
# planeMasks_expanded = tf.expand_dims(planeMasks, 0)
148+
# boundary = tf.reduce_max(tf.nn.max_pool(planeMasks_expanded, ksize=[1, 3, 3, 1], strides=[1, 1, 1, 1], padding='SAME', name='max_pool') - planeMasks_expanded, axis=3, keep_dims=True)
149+
# max_depth_diff = 0.1
150+
# depth_expanded = tf.expand_dims(depth, 0)
151+
# kernel_size = 5
152+
# padding = (kernel_size - 1) / 2
153+
# neighbor_kernel_array = gaussian(kernel_size, kernel_size)
154+
# neighbor_kernel_array[(kernel_size - 1) / 2][(kernel_size - 1) / 2] = 0
155+
# neighbor_kernel_array /= neighbor_kernel_array.sum()
156+
# neighbor_kernel_array *= -1
157+
# neighbor_kernel_array[(kernel_size - 1) / 2][(kernel_size - 1) / 2] = 1
158+
# neighbor_kernel = tf.constant(neighbor_kernel_array.reshape(-1), shape=neighbor_kernel_array.shape, dtype=tf.float32)
159+
# neighbor_kernel = tf.reshape(neighbor_kernel, [kernel_size, kernel_size, 1, 1])
160+
161+
# depth_diff = tf.abs(tf.nn.depthwise_conv2d(depth_expanded, neighbor_kernel, strides=[1, 1, 1, 1], padding='VALID'))
162+
# depth_diff = tf.pad(depth_diff, paddings = [[0, 0], [padding, padding], [padding, padding], [0, 0]])
163+
# smooth_boundary = boundary * tf.cast(tf.less(depth_diff, max_depth_diff), tf.float32)
164+
# occlusion_boundary = boundary - smooth_boundary
165+
# boundary = tf.squeeze(tf.concat([smooth_boundary, occlusion_boundary], axis=3), axis=0)
166+
167+
168+
#validating = tf.cast(features['validating'], tf.float32)
169+
shuffle_inds = tf.one_hot(tf.random_shuffle(tf.range(numPlanes)), depth = numPlanes)
170+
#shuffle_inds = tf.one_hot(tf.range(numPlanes), depth = numPlanes)
171+
172+
#shuffle_inds = tf.concat([shuffle_inds, tf.zeros((numPlanes, numOutputPlanes - numPlanes))], axis=1)
173+
#shuffle_inds = tf.concat([shuffle_inds, tf.concat([tf.zeros((numOutputPlanes - numPlanes, numPlanes)), tf.diag(tf.ones([numOutputPlanes - numPlanes]))], axis=1)], axis=0)
174+
planes = tf.transpose(tf.matmul(tf.transpose(planes), shuffle_inds))
175+
planes = tf.reshape(planes, [numPlanes, 3])
176+
planes = tf.concat([planes, tf.zeros([numOutputPlanes - numPlanes, 3])], axis=0)
177+
planes = tf.reshape(planes, [numOutputPlanes, 3])
178+
179+
180+
boundary = tf.decode_raw(features['boundary_raw'], tf.uint8)
181+
boundary = tf.cast(boundary > 128, tf.float32)
182+
boundary = tf.reshape(boundary, [HEIGHT, WIDTH, 2])
183+
#boundary = tf.slice(tf.reshape(boundary, [HEIGHT, WIDTH, 3]), [0, 0, 0], [HEIGHT, WIDTH, 2])
184+
185+
grid_s = tf.reshape(features['grid_s'], [HEIGHT / 8, WIDTH / 8, 1])
186+
grid_p = tf.reshape(features['grid_p'], [HEIGHT / 8, WIDTH / 8, 3])
187+
188+
grid_m = tf.decode_raw(features['grid_m_raw'], tf.uint8)
189+
grid_m = tf.cast(tf.reshape(grid_m, [HEIGHT / 8, WIDTH / 8, 16 * 16]), tf.float32)
190+
191+
if getLocal:
192+
if random:
193+
image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp, boundary_gt, grid_s_gt, grid_p_gt, grid_m_gt, num_planes_gt = tf.train.shuffle_batch([image, planes, depth, normal, planeMasks, boundary, grid_s, grid_p, grid_m, numPlanes], batch_size=batchSize, capacity=min_after_dequeue + (NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS, min_after_dequeue=min_after_dequeue)
194+
else:
195+
image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp, boundary_gt, grid_s_gt, grid_p_gt, grid_m_gt, num_planes_gt = tf.train.batch([image, planes, depth, normal, planeMasks, boundary, grid_s, grid_p, grid_m, numPlanes], batch_size=batchSize, capacity=(NUM_THREADS + 2) * batchSize, num_threads=1)
196+
pass
197+
return image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp, boundary_gt, grid_s_gt, grid_p_gt, grid_m_gt, num_planes_gt
198+
199+
200+
if not getSegmentation:
201+
if random:
202+
image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp = tf.train.shuffle_batch([image, planes, depth, normal, tf.zeros([HEIGHT, WIDTH, numOutputPlanes])], batch_size=batchSize, capacity=min_after_dequeue + (NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS, min_after_dequeue=min_after_dequeue)
203+
else:
204+
image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp = tf.train.batch([image, planes, depth, normal, tf.zeros([HEIGHT, WIDTH, numOutputPlanes])], batch_size=batchSize, capacity=(NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS)
205+
pass
206+
return image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp
207+
208+
209+
plane_masks = tf.cast(features['plane_mask'], tf.int64)
210+
211+
plane_masks = tf.reshape(plane_masks, [HEIGHT, WIDTH, 1])
212+
plane_mask_array = tf.tile(plane_masks, [1, 1, numPlanes])
213+
coef = tf.range(numPlanes, dtype=tf.int64)
214+
coef = tf.pow(2, coef)
215+
coef = tf.reshape(tf.matmul(tf.reshape(coef, [-1, numPlanes]), tf.cast(shuffle_inds, tf.int64)), [numPlanes])
216+
coef = tf.cast(coef, tf.int64)
217+
plane_mask_array = tf.cast(tf.div(plane_mask_array, coef) % 2, tf.float32)
218+
plane_mask_array = tf.concat([plane_mask_array, tf.zeros([HEIGHT, WIDTH, numOutputPlanes - numPlanes])], axis=2)
219+
plane_mask_array = tf.reshape(plane_mask_array, [HEIGHT, WIDTH, numOutputPlanes])
220+
#num_planes_array = tf.concat([tf.ones([numPlanes], dtype=np.float32) / tf.cast(numPlanes * BATCH_SIZE, np.float32), tf.zeros([numOutputPlanes - numPlanes])], axis=0)
221+
222+
#image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp, num_planes, mask = tf.train.shuffle_batch([image, planes, depth, normal, plane_mask_array, numPlanes, plane_masks_test], batch_size=batchSize, capacity=min_after_dequeue + (NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS, min_after_dequeue=min_after_dequeue)
223+
224+
# if True:
225+
# image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp = tf.train.batch([image, planes, tf.ones((HEIGHT, WIDTH, 1)), tf.ones((HEIGHT, WIDTH, 3)), plane_mask_array], batch_size=batchSize, capacity=(NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS)
226+
# return image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp
227+
228+
if random:
229+
image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp, boundary_gt, grid_s_gt, grid_p_gt, grid_m_gt, num_planes_gt = tf.train.shuffle_batch([image, planes, depth, normal, plane_mask_array, boundary, grid_s, grid_p, grid_m, numPlanes], batch_size=batchSize, capacity=min_after_dequeue + (NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS, min_after_dequeue=min_after_dequeue)
230+
image_inp, plane_inp, depth_gt, normal_gt, plane_mask_gt, boundary_gt, num_planes_gt = tf.train.shuffle_batch([image, planes, depth, normal, plane_mask_array, boundary, numPlanes], batch_size=batchSize, capacity=min_after_dequeue + (NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS, min_after_dequeue=min_after_dequeue)
231+
else:
232+
image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp, boundary_gt, grid_s_gt, grid_p_gt, grid_m_gt, num_planes_gt = tf.train.batch([image, planes, depth, normal, plane_mask_array, boundary, grid_s, grid_p, grid_m, numPlanes], batch_size=batchSize, capacity=(NUM_THREADS + 2) * batchSize, num_threads=1)
233+
pass
234+
return image_inp, plane_inp, depth_gt, normal_gt, plane_mask_inp, boundary_gt, grid_s_gt, grid_p_gt, grid_m_gt, num_planes_gt

‎PlaneSetGeneration/layers.py

+57
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import os
2+
import numpy as np
3+
4+
#DEBUG = False
5+
6+
class RangesLayer(object):
7+
def __init__(self, width, height):
8+
9+
focalLength = 517.97
10+
urange = np.arange(width).reshape(1, -1).repeat(height, 0) - width * 0.5
11+
vrange = np.arange(height).reshape(-1, 1).repeat(width, 1) - height * 0.5
12+
self.ranges = np.array([urange / focalLength / width * 640, np.ones(urange.shape), -vrange / focalLength / height * 480]).transpose([1, 2, 0])
13+
return
14+
15+
def forward(self):
16+
return self.ranges
17+
18+
19+
def PlaneDepthLayer(planes, ranges):
20+
batchSize = 1
21+
if len(planes.shape) == 3:
22+
batchSize = planes.shape[0]
23+
planes = planes.reshape(planes.shape[0] * planes.shape[1], planes.shape[2])
24+
pass
25+
26+
planesD = np.linalg.norm(planes, 2, 1)
27+
planesD = np.maximum(planesD, 1e-4)
28+
planesNormal = -planes / planesD.reshape(-1, 1).repeat(3, 1)
29+
30+
normalXYZ = np.dot(ranges, planesNormal.transpose())
31+
normalXYZ[normalXYZ == 0] = 1e-4
32+
normalXYZ = 1 / normalXYZ
33+
depths = -normalXYZ
34+
depths[:, :] *= planesD
35+
if batchSize > 1:
36+
depths = depths.reshape(depths.shape[0], depths.shape[1], batchSize, -1).transpose([2, 0, 1, 3])
37+
pass
38+
depths[(depths < 0) + (depths > 10)] = 10
39+
#depths[depths < 0] = 0
40+
#depths[depths > 10] = 10
41+
return depths
42+
43+
44+
def PlaneNormalLayer(planes, ranges):
45+
batchSize = 1
46+
if len(planes.shape) == 3:
47+
batchSize = planes.shape[0]
48+
planes = planes.reshape(planes.shape[0] * planes.shape[1], planes.shape[2])
49+
pass
50+
planesD = np.linalg.norm(planes, 2, 1)
51+
planesD = np.maximum(planesD, 1e-4)
52+
planesNormal = -planes / planesD.reshape(-1, 1).repeat(3, 1)
53+
normals = planesNormal.reshape(1, 1, -1, 3).repeat(ranges.shape[0], 0).repeat(ranges.shape[1], 1)
54+
if batchSize > 1:
55+
normals = normals.reshape(normals.shape[0], normals.shape[1], batchSize, -1, 3).transpose([2, 0, 1, 3, 4])
56+
pass
57+
return normals

0 commit comments

Comments
 (0)