Commit db1caa42 authored by Yuxin Wu's avatar Yuxin Wu

fix PredictConfig. use inputs_desc adn tower_func for load-cpm

parent bbb2ecc2
......@@ -44,71 +44,65 @@ def get_gaussian_map():
return gaussian_map.reshape((1, 368, 368, 1))
class Model(ModelDesc):
def _get_inputs(self):
return [InputDesc(tf.float32, (None, 368, 368, 3), 'input'),
InputDesc(tf.float32, (None, 368, 368, 15), 'label'),
]
def _build_graph(self, inputs):
image, label = inputs
image = image / 256.0 - 0.5
gmap = tf.constant(get_gaussian_map())
gmap = tf.pad(gmap, [[0, 0], [0, 1], [0, 1], [0, 0]])
pool_center = AvgPooling('mappool', gmap, 9, stride=8, padding='VALID')
with argscope(Conv2D, kernel_shape=3, nl=tf.nn.relu,
W_init=tf.random_normal_initializer(stddev=0.01)):
shared = (LinearWrap(image)
.Conv2D('conv1_1', 64)
.Conv2D('conv1_2', 64)
.MaxPooling('pool1', 2)
# 184
.Conv2D('conv2_1', 128)
.Conv2D('conv2_2', 128)
.MaxPooling('pool2', 2)
# 92
.Conv2D('conv3_1', 256)
.Conv2D('conv3_2', 256)
.Conv2D('conv3_3', 256)
.Conv2D('conv3_4', 256)
.MaxPooling('pool3', 2)
# 46
.Conv2D('conv4_1', 512)
.Conv2D('conv4_2', 512)
.Conv2D('conv4_3_CPM', 256)
.Conv2D('conv4_4_CPM', 256)
.Conv2D('conv4_5_CPM', 256)
.Conv2D('conv4_6_CPM', 256)
.Conv2D('conv4_7_CPM', 128)())
def add_stage(stage, l):
l = tf.concat([l, shared, pool_center], 3,
name='concat_stage{}'.format(stage))
for i in range(1, 6):
l = Conv2D('Mconv{}_stage{}'.format(i, stage), l, 128)
l = Conv2D('Mconv6_stage{}'.format(stage), l, 128, kernel_shape=1)
l = Conv2D('Mconv7_stage{}'.format(stage),
l, 15, kernel_shape=1, nl=tf.identity)
return l
with argscope(Conv2D, kernel_shape=7, nl=tf.nn.relu):
out1 = (LinearWrap(shared)
.Conv2D('conv5_1_CPM', 512, kernel_shape=1)
.Conv2D('conv5_2_CPM', 15, kernel_shape=1, nl=tf.identity)())
out2 = add_stage(2, out1)
out3 = add_stage(3, out2)
out4 = add_stage(4, out3)
out5 = add_stage(5, out4)
out6 = add_stage(6, out4)
resized_map = tf.image.resize_bilinear(out6,
[368, 368], name='resized_map')
def CPM(image):
image = image / 256.0 - 0.5
gmap = tf.constant(get_gaussian_map())
gmap = tf.pad(gmap, [[0, 0], [0, 1], [0, 1], [0, 0]])
pool_center = AvgPooling('mappool', gmap, 9, stride=8, padding='VALID')
with argscope(Conv2D, kernel_shape=3, nl=tf.nn.relu,
W_init=tf.random_normal_initializer(stddev=0.01)):
shared = (LinearWrap(image)
.Conv2D('conv1_1', 64)
.Conv2D('conv1_2', 64)
.MaxPooling('pool1', 2)
# 184
.Conv2D('conv2_1', 128)
.Conv2D('conv2_2', 128)
.MaxPooling('pool2', 2)
# 92
.Conv2D('conv3_1', 256)
.Conv2D('conv3_2', 256)
.Conv2D('conv3_3', 256)
.Conv2D('conv3_4', 256)
.MaxPooling('pool3', 2)
# 46
.Conv2D('conv4_1', 512)
.Conv2D('conv4_2', 512)
.Conv2D('conv4_3_CPM', 256)
.Conv2D('conv4_4_CPM', 256)
.Conv2D('conv4_5_CPM', 256)
.Conv2D('conv4_6_CPM', 256)
.Conv2D('conv4_7_CPM', 128)())
def add_stage(stage, l):
l = tf.concat([l, shared, pool_center], 3,
name='concat_stage{}'.format(stage))
for i in range(1, 6):
l = Conv2D('Mconv{}_stage{}'.format(i, stage), l, 128)
l = Conv2D('Mconv6_stage{}'.format(stage), l, 128, kernel_shape=1)
l = Conv2D('Mconv7_stage{}'.format(stage),
l, 15, kernel_shape=1, nl=tf.identity)
return l
with argscope(Conv2D, kernel_shape=7, nl=tf.nn.relu):
out1 = (LinearWrap(shared)
.Conv2D('conv5_1_CPM', 512, kernel_shape=1)
.Conv2D('conv5_2_CPM', 15, kernel_shape=1, nl=tf.identity)())
out2 = add_stage(2, out1)
out3 = add_stage(3, out2)
out4 = add_stage(4, out3)
out5 = add_stage(5, out4)
out6 = add_stage(6, out4)
resized_map = tf.image.resize_bilinear(out6,
[368, 368], name='resized_map')
def run_test(model_path, img_file):
param_dict = np.load(model_path, encoding='latin1').item()
predict_func = OfflinePredictor(PredictConfig(
model=Model(),
inputs_desc=[InputDesc(tf.float32, (None, 368, 368, 3), 'input')],
tower_func=CPM,
session_init=DictRestore(param_dict),
input_names=['input'],
output_names=['resized_map']
......
......@@ -27,7 +27,7 @@ class PredictConfig(object):
):
"""
Args:
model (ModelDescBase): the model to be used to obtain inputs_desc and tower_func.
model (ModelDescBase): to be used to obtain inputs_desc and tower_func.
inputs_desc ([InputDesc]):
tower_func: a callable which takes input tensors
......@@ -35,8 +35,7 @@ class PredictConfig(object):
session. Defaults to :class:`tf.train.ChiefSessionCreator()`.
session_init (SessionInit): how to initialize variables of the session.
Defaults to do nothing.
input_names (list): a list of input tensor names. Defaults to all
inputs of the model.
input_names (list): a list of input tensor names. Defaults to match inputs_desc.
output_names (list): a list of names of the output tensors to predict, the
tensors can be any computable tensor in the graph.
return_input (bool): same as in :attr:`PredictorBase.return_input`.
......@@ -70,8 +69,7 @@ class PredictConfig(object):
# inputs & outputs
self.input_names = input_names
if self.input_names is None:
raw_tensors = self.model.get_inputs_desc()
self.input_names = [k.name for k in raw_tensors]
self.input_names = [k.name for k in self.inputs_desc]
self.output_names = output_names
assert_type(self.output_names, list)
assert_type(self.input_names, list)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment