Skip to content

Commit e019207

Browse files
art-programmerart-programmer
committed
minor changes
0 parents  commit e019207

12 files changed

+16505
-0
lines changed

‎RecordReaderAll.py

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
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+
HEIGHT=192
17+
WIDTH=256
18+
NUM_PLANES = 20
19+
NUM_THREADS = 4
20+
21+
22+
23+
class RecordReaderAll():
24+
def __init__(self):
25+
return
26+
27+
def getBatch(self, filename_queue, numOutputPlanes = 20, batchSize = 16, min_after_dequeue = 1000, random=True, getLocal=False, getSegmentation=False, test=True):
28+
reader = tf.TFRecordReader()
29+
_, serialized_example = reader.read(filename_queue)
30+
31+
features = tf.parse_single_example(
32+
serialized_example,
33+
# Defaults are not specified since both keys are required.
34+
features={
35+
#'height': tf.FixedLenFeature([], tf.int64),
36+
#'width': tf.FixedLenFeature([], tf.int64),
37+
'image_raw': tf.FixedLenFeature([], tf.string),
38+
'image_path': tf.FixedLenFeature([], tf.string),
39+
'num_planes': tf.FixedLenFeature([], tf.int64),
40+
'plane': tf.FixedLenFeature([NUM_PLANES * 3], tf.float32),
41+
#'plane_relation': tf.FixedLenFeature([NUM_PLANES * NUM_PLANES], tf.float32),
42+
'segmentation_raw': tf.FixedLenFeature([], tf.string),
43+
'depth': tf.FixedLenFeature([HEIGHT * WIDTH], tf.float32),
44+
'normal': tf.FixedLenFeature([HEIGHT * WIDTH * 3], tf.float32),
45+
'semantics_raw': tf.FixedLenFeature([], tf.string),
46+
'boundary_raw': tf.FixedLenFeature([], tf.string),
47+
'info': tf.FixedLenFeature([4 * 4 + 4], tf.float32),
48+
})
49+
50+
# Convert from a scalar string tensor (whose single string has
51+
# length mnist.IMAGE_PIXELS) to a uint8 tensor with shape
52+
# [mnist.IMAGE_PIXELS].
53+
image = tf.decode_raw(features['image_raw'], tf.uint8)
54+
image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
55+
image = tf.reshape(image, [HEIGHT, WIDTH, 3])
56+
57+
58+
depth = features['depth']
59+
depth = tf.reshape(depth, [HEIGHT, WIDTH, 1])
60+
61+
normal = features['normal']
62+
normal = tf.reshape(normal, [HEIGHT, WIDTH, 3])
63+
64+
normal = tf.nn.l2_normalize(normal, dim=2)
65+
66+
#normal = tf.stack([normal[:, :, 1], normal[:, :, 0], normal[:, :, 2]], axis=2)
67+
68+
69+
semantics = tf.decode_raw(features['semantics_raw'], tf.uint8)
70+
semantics = tf.cast(tf.reshape(semantics, [HEIGHT, WIDTH]), tf.int32)
71+
72+
numPlanes = tf.minimum(tf.cast(features['num_planes'], tf.int32), numOutputPlanes)
73+
74+
numPlanesOri = numPlanes
75+
numPlanes = tf.maximum(numPlanes, 1)
76+
77+
planes = features['plane']
78+
planes = tf.reshape(planes, [NUM_PLANES, 3])
79+
planes = tf.slice(planes, [0, 0], [numPlanes, 3])
80+
81+
#shuffle_inds = tf.one_hot(tf.random_shuffle(tf.range(numPlanes)), depth = numPlanes)
82+
shuffle_inds = tf.one_hot(tf.range(numPlanes), numPlanes)
83+
84+
planes = tf.transpose(tf.matmul(tf.transpose(planes), shuffle_inds))
85+
planes = tf.reshape(planes, [numPlanes, 3])
86+
planes = tf.concat([planes, tf.zeros([numOutputPlanes - numPlanes, 3])], axis=0)
87+
planes = tf.reshape(planes, [numOutputPlanes, 3])
88+
89+
90+
boundary = tf.decode_raw(features['boundary_raw'], tf.uint8)
91+
boundary = tf.cast(tf.reshape(boundary, (HEIGHT, WIDTH, 2)), tf.float32)
92+
93+
#boundary = tf.decode_raw(features['boundary_raw'], tf.float64)
94+
#boundary = tf.cast(tf.reshape(boundary, (HEIGHT, WIDTH, 3)), tf.float32)
95+
#boundary = tf.slice(boundary, [0, 0, 0], [HEIGHT, WIDTH, 2])
96+
97+
segmentation = tf.decode_raw(features['segmentation_raw'], tf.uint8)
98+
segmentation = tf.reshape(segmentation, [HEIGHT, WIDTH, 1])
99+
100+
101+
102+
coef = tf.range(numPlanes)
103+
coef = tf.reshape(tf.matmul(tf.reshape(coef, [-1, numPlanes]), tf.cast(shuffle_inds, tf.int32)), [1, 1, numPlanes])
104+
105+
plane_masks = tf.cast(tf.equal(segmentation, tf.cast(coef, tf.uint8)), tf.float32)
106+
plane_masks = tf.concat([plane_masks, tf.zeros([HEIGHT, WIDTH, numOutputPlanes - numPlanes])], axis=2)
107+
plane_masks = tf.reshape(plane_masks, [HEIGHT, WIDTH, numOutputPlanes])
108+
109+
#non_plane_mask = tf.cast(tf.equal(segmentation, tf.cast(numOutputPlanes, tf.uint8)), tf.float32)
110+
non_plane_mask = 1 - tf.reduce_max(plane_masks, axis=2, keep_dims=True)
111+
#tf.cast(tf.equal(segmentation, tf.cast(numOutputPlanes, tf.uint8)), tf.float32)
112+
113+
114+
if random:
115+
image_inp, plane_inp, depth_gt, normal_gt, semantics_gt, plane_masks_gt, boundary_gt, num_planes_gt, non_plane_mask_gt, image_path, info = tf.train.shuffle_batch([image, planes, depth, normal, semantics, plane_masks, boundary, numPlanesOri, non_plane_mask, features['image_path'], features['info']], batch_size=batchSize, capacity=min_after_dequeue + (NUM_THREADS + 2) * batchSize, num_threads=NUM_THREADS, min_after_dequeue=min_after_dequeue)
116+
else:
117+
image_inp, plane_inp, depth_gt, normal_gt, semantics_gt, plane_masks_gt, boundary_gt, num_planes_gt, non_plane_mask_gt, image_path, info = tf.train.batch([image, planes, depth, normal, semantics, plane_masks, boundary, numPlanesOri, non_plane_mask, features['image_path'], features['info']], batch_size=batchSize, capacity=(NUM_THREADS + 2) * batchSize, num_threads=1)
118+
pass
119+
global_gt_dict = {'plane': plane_inp, 'depth': depth_gt, 'normal': normal_gt, 'semantics': semantics_gt, 'segmentation': plane_masks_gt, 'boundary': boundary_gt, 'num_planes': num_planes_gt, 'non_plane_mask': non_plane_mask_gt, 'image_path': image_path, 'info': info}
120+
return image_inp, global_gt_dict, {}

‎crfasrnn_layer.py

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
"""
2+
MIT License
3+
Copyright (c) 2017 Sadeep Jayasumana
4+
Permission is hereby granted, free of charge, to any person obtaining a copy
5+
of this software and associated documentation files (the "Software"), to deal
6+
in the Software without restriction, including without limitation the rights
7+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
8+
copies of the Software, and to permit persons to whom the Software is
9+
furnished to do so, subject to the following conditions:
10+
The above copyright notice and this permission notice shall be included in all
11+
copies or substantial portions of the Software.
12+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
13+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
14+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
15+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
16+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
17+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
18+
SOFTWARE.
19+
"""
20+
21+
import numpy as np
22+
import tensorflow as tf
23+
from keras.engine.topology import Layer
24+
custom_module = tf.load_op_library('./cpp/high_dim_filter.so')
25+
import high_dim_filter_grad # Register gradients for the custom op
26+
27+
28+
class CrfRnnLayer(Layer):
29+
""" Implements the CRF-RNN layer described in:
30+
Conditional Random Fields as Recurrent Neural Networks,
31+
S. Zheng, S. Jayasumana, B. Romera-Paredes, V. Vineet, Z. Su, D. Du, C. Huang and P. Torr,
32+
ICCV 2015
33+
"""
34+
35+
def __init__(self, image_dims, num_classes,
36+
theta_alpha, theta_beta, theta_gamma,
37+
num_iterations, **kwargs):
38+
self.image_dims = image_dims
39+
self.num_classes = num_classes
40+
self.theta_alpha = theta_alpha
41+
self.theta_beta = theta_beta
42+
self.theta_gamma = theta_gamma
43+
self.num_iterations = num_iterations
44+
self.spatial_ker_weights = None
45+
self.bilateral_ker_weights = None
46+
self.compatibility_matrix = None
47+
super(CrfRnnLayer, self).__init__(**kwargs)
48+
49+
def build(self, input_shape):
50+
# # Weights of the spatial kernel
51+
# self.spatial_ker_weights = self.add_weight(name='spatial_ker_weights',
52+
# shape=(self.num_classes, self.num_classes),
53+
# initializer='uniform',
54+
# trainable=True)
55+
56+
# # Weights of the bilateral kernel
57+
# self.bilateral_ker_weights = self.add_weight(name='bilateral_ker_weights',
58+
# shape=(self.num_classes, self.num_classes),
59+
# initializer='uniform',
60+
# trainable=True)
61+
62+
# # Compatibility matrix
63+
# self.compatibility_matrix = self.add_weight(name='compatibility_matrix',
64+
# shape=(self.num_classes, self.num_classes),
65+
# initializer='uniform',
66+
# trainable=True)
67+
68+
69+
weights = np.load('weights.npy')
70+
weights = [weights[0], weights[1], weights[2]]
71+
self.spatial_ker_weights = tf.Variable(weights[0][:self.num_classes, :self.num_classes], name='spatial_ker_weights', trainable=True)
72+
self.bilateral_ker_weights = tf.Variable(weights[1][:self.num_classes, :self.num_classes], name='bilateral_ker_weights', trainable=True)
73+
self.compatibility_matrix = tf.Variable(weights[2][:self.num_classes, :self.num_classes], name='compatibility_matrix', trainable=True)
74+
75+
76+
# self.spatial_ker_weights = tf.constant(weights[0].reshape(-1), name='spatial_ker_weights', shape=(self.num_classes, self.num_classes))
77+
# self.bilateral_ker_weights = tf.constant(weights[1].reshape(-1), name='bilateral_ker_weights', shape=(self.num_classes, self.num_classes))
78+
# self.compatibility_ker_weights = tf.constant(weights[2].reshape(-1), name='compatibility_ker_weights', shape=(self.num_classes, self.num_classes))
79+
80+
81+
super(CrfRnnLayer, self).build(input_shape)
82+
83+
84+
def call(self, inputs):
85+
batchSize = int(inputs[0].shape[0])
86+
c, h, w = self.num_classes, self.image_dims[0], self.image_dims[1]
87+
all_ones = np.ones((c, h, w), dtype=np.float32)
88+
89+
outputs = []
90+
for batchIndex in xrange(batchSize):
91+
unaries = tf.transpose(inputs[0][batchIndex, :, :, :], perm=(2, 0, 1))
92+
rgb = tf.transpose(inputs[1][batchIndex, :, :, :], perm=(2, 0, 1))
93+
94+
95+
# Prepare filter normalization coefficients
96+
spatial_norm_vals = custom_module.high_dim_filter(all_ones, rgb, bilateral=False,
97+
theta_gamma=self.theta_gamma)
98+
bilateral_norm_vals = custom_module.high_dim_filter(all_ones, rgb, bilateral=True,
99+
theta_alpha=self.theta_alpha,
100+
theta_beta=self.theta_beta)
101+
q_values = unaries
102+
103+
for i in range(self.num_iterations):
104+
softmax_out = tf.nn.softmax(q_values, dim=0)
105+
106+
# Spatial filtering
107+
spatial_out = custom_module.high_dim_filter(softmax_out, rgb, bilateral=False,
108+
theta_gamma=self.theta_gamma)
109+
spatial_out = spatial_out / spatial_norm_vals
110+
111+
# Bilateral filtering
112+
bilateral_out = custom_module.high_dim_filter(softmax_out, rgb, bilateral=True,
113+
theta_alpha=self.theta_alpha,
114+
theta_beta=self.theta_beta)
115+
bilateral_out = bilateral_out / bilateral_norm_vals
116+
117+
# Weighting filter outputs
118+
message_passing = (tf.matmul(self.spatial_ker_weights,
119+
tf.reshape(spatial_out, (c, -1))) +
120+
tf.matmul(self.bilateral_ker_weights,
121+
tf.reshape(bilateral_out, (c, -1))))
122+
123+
# Compatibility transform
124+
pairwise = tf.matmul(self.compatibility_matrix, message_passing)
125+
126+
# Adding unary potentials
127+
pairwise = tf.reshape(pairwise, (c, h, w))
128+
q_values = unaries - pairwise
129+
continue
130+
outputs.append(tf.transpose(tf.reshape(q_values, (1, c, h, w)), perm=(0, 2, 3, 1)))
131+
continue
132+
outputs = tf.concat(outputs, axis=0)
133+
return outputs
134+
135+
def compute_output_shape(self, input_shape):
136+
return input_shape

0 commit comments

Comments
 (0)