Skip to content

Commit 26e2b09

Browse files
zsdonghaozsdonghao
committed
yapf
1 parent fe46e14 commit 26e2b09

File tree

6 files changed

+106
-206
lines changed

6 files changed

+106
-206
lines changed

‎config.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
config.TRAIN.gamma = 0.333 # gamma of Adam
1515
config.TRAIN.weight_decay = 5e-4
1616
config.TRAIN.distributed = False
17-
config.TRAIN.train_mode == 'placeholder' # placeholder, datasetapi, distributed
17+
config.TRAIN.train_mode == 'placeholder' # placeholder, datasetapi, distributed
1818

1919
config.MODEL = edict()
2020
config.MODEL.model_path = 'models' # save directory

‎inference.py

Lines changed: 10 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -31,23 +31,16 @@
3131
_, _, _, net = model(x, n_pos, None, None, False, False)
3232

3333
# get output from network
34-
conf_tensor = tl.layers.get_layers_with_name(
35-
net, 'model/cpm/stage6/branch1/conf')[0]
36-
pafs_tensor = tl.layers.get_layers_with_name(
37-
net, 'model/cpm/stage6/branch2/pafs')[0]
34+
conf_tensor = tl.layers.get_layers_with_name(net, 'model/cpm/stage6/branch1/conf')[0]
35+
pafs_tensor = tl.layers.get_layers_with_name(net, 'model/cpm/stage6/branch2/pafs')[0]
3836

3937
def get_peak(pafs_tensor):
4038
from inference.smoother import Smoother
4139
smoother = Smoother({'data': pafs_tensor}, 25, 3.0)
4240
gaussian_heatMat = smoother.get_output()
43-
max_pooled_in_tensor = tf.nn.pool(
44-
gaussian_heatMat,
45-
window_shape=(3, 3),
46-
pooling_type='MAX',
47-
padding='SAME')
41+
max_pooled_in_tensor = tf.nn.pool(gaussian_heatMat, window_shape=(3, 3), pooling_type='MAX', padding='SAME')
4842
tensor_peaks = tf.where(
49-
tf.equal(gaussian_heatMat, max_pooled_in_tensor), gaussian_heatMat,
50-
tf.zeros_like(gaussian_heatMat))
43+
tf.equal(gaussian_heatMat, max_pooled_in_tensor), gaussian_heatMat, tf.zeros_like(gaussian_heatMat))
5144
return tensor_peaks
5245

5346
peak_tensor = get_peak(pafs_tensor)
@@ -56,8 +49,7 @@ def get_peak(pafs_tensor):
5649
sess = tf.InteractiveSession()
5750
sess.run(tf.global_variables_initializer())
5851
if model_file:
59-
tl.files.load_and_assign_npz_dict(
60-
os.path.join(model_path, model_file), sess)
52+
tl.files.load_and_assign_npz_dict(os.path.join(model_path, model_file), sess)
6153

6254
# get one example image with range 0~1
6355
im = tl.vis.read_image(input_file)
@@ -68,8 +60,7 @@ def get_peak(pafs_tensor):
6860
# 1st time need time to compile
6961
# _, _ = sess.run([conf_tensor, pafs_tensor], feed_dict={x: [im]})
7062
st = time.time()
71-
conf, pafs, peak = sess.run(
72-
[conf_tensor, pafs_tensor, peak_tensor], feed_dict={x: [im]})
63+
conf, pafs, peak = sess.run([conf_tensor, pafs_tensor, peak_tensor], feed_dict={x: [im]})
7364
t = time.time() - st
7465
print("get maps took {}s i.e. {} FPS".format(t, 1. / t))
7566
# print(conf.shape, pafs.shape, peak.shape)
@@ -92,11 +83,10 @@ def estimate_paf(peaks, heat_mat, paf_mat):
9283
continue
9384

9485
is_added = True
95-
human.body_parts[part_idx] = BodyPart(
96-
'%d-%d' % (human_id, part_idx), part_idx,
97-
float(pafprocess.get_part_x(c_idx)) / heat_mat.shape[1],
98-
float(pafprocess.get_part_y(c_idx)) / heat_mat.shape[0],
99-
pafprocess.get_part_score(c_idx))
86+
human.body_parts[part_idx] = BodyPart('%d-%d' % (human_id, part_idx), part_idx,
87+
float(pafprocess.get_part_x(c_idx)) / heat_mat.shape[1],
88+
float(pafprocess.get_part_y(c_idx)) / heat_mat.shape[0],
89+
pafprocess.get_part_score(c_idx))
10090

10191
if is_added:
10292
score = pafprocess.get_score(human_id)

‎models.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ def stage(cnn, b1, b2, n_pos, maskInput1, maskInput2, is_train, name='stageX'):
2929
b2 = Conv2d(b2, 128, (7, 7), (1, 1), tf.nn.relu, 'SAME', W_init=W_init, b_init=b_init, name='c4')
3030
b2 = Conv2d(b2, 128, (7, 7), (1, 1), tf.nn.relu, 'SAME', W_init=W_init, b_init=b_init, name='c5')
3131
b2 = Conv2d(b2, 128, (1, 1), (1, 1), tf.nn.relu, 'VALID', W_init=W_init, b_init=b_init, name='c6')
32-
b2 = Conv2d(b2, 38, (1, 1), (1, 1), None, 'VALID', W_init=W_init, b_init=b_init, name='pafs')
32+
b2 = Conv2d(b2, n_pos * 2, (1, 1), (1, 1), None, 'VALID', W_init=W_init, b_init=b_init, name='pafs')
3333
if is_train:
3434
b2.outputs = b2.outputs * maskInput2
3535
return b1, b2
@@ -87,25 +87,26 @@ def model(x, n_pos, mask_miss1, mask_miss2, is_train=False, reuse=None):
8787
b1 = Conv2d(cnn, 128, (3, 3), (1, 1), tf.nn.relu, 'SAME', W_init=W_init, b_init=b_init, name='c1')
8888
b1 = Conv2d(b1, 128, (3, 3), (1, 1), tf.nn.relu, 'SAME', W_init=W_init, b_init=b_init, name='c2')
8989
b1 = Conv2d(b1, 128, (3, 3), (1, 1), tf.nn.relu, 'SAME', W_init=W_init, b_init=b_init, name='c3')
90-
b1 = Conv2d(b1, 512, (1, 1), (1, 1), tf.nn.relu, 'VALID', W_init=W_init, b_init=b_init, name='c4')
90+
# b1 = Conv2d(b1, 512, (1, 1), (1, 1), tf.nn.relu, 'VALID', W_init=W_init, b_init=b_init, name='c4')
91+
b1 = Conv2d(b1, 128, (1, 1), (1, 1), tf.nn.relu, 'VALID', W_init=W_init, b_init=b_init, name='c4')
9192
b1 = Conv2d(b1, n_pos, (1, 1), (1, 1), None, 'VALID', W_init=W_init, b_init=b_init, name='confs')
9293
if is_train:
9394
b1.outputs = b1.outputs * mask_miss1
9495
with tf.variable_scope("stage1/branch2"):
9596
b2 = Conv2d(cnn, 128, (3, 3), (1, 1), tf.nn.relu, 'SAME', W_init=W_init, b_init=b_init, name='c1')
9697
b2 = Conv2d(b2, 128, (3, 3), (1, 1), tf.nn.relu, 'SAME', W_init=W_init, b_init=b_init, name='c2')
9798
b2 = Conv2d(b2, 128, (3, 3), (1, 1), tf.nn.relu, 'SAME', W_init=W_init, b_init=b_init, name='c3')
98-
b2 = Conv2d(b2, 512, (1, 1), (1, 1), tf.nn.relu, 'VALID', W_init=W_init, b_init=b_init, name='c4')
99-
b2 = Conv2d(b2, 38, (1, 1), (1, 1), None, 'VALID', W_init=W_init, b_init=b_init, name='pafs')
99+
# b2 = Conv2d(b2, 512, (1, 1), (1, 1), tf.nn.relu, 'VALID', W_init=W_init, b_init=b_init, name='c4')
100+
b2 = Conv2d(b2, 128, (1, 1), (1, 1), tf.nn.relu, 'VALID', W_init=W_init, b_init=b_init, name='c4')
101+
b2 = Conv2d(b2, n_pos * 2, (1, 1), (1, 1), None, 'VALID', W_init=W_init, b_init=b_init, name='pafs')
100102
if is_train:
101103
b2.outputs = b2.outputs * mask_miss2
102104
b1_list.append(b1)
103105
b2_list.append(b2)
104106
# stage 2~6
105107
for i in range(2, 7):
106108
b1, b2 = stage(
107-
cnn, b1_list[-1], b2_list[-1], n_pos, mask_miss1, mask_miss2, is_train, name='stage%d' % i
108-
)
109+
cnn, b1_list[-1], b2_list[-1], n_pos, mask_miss1, mask_miss2, is_train, name='stage%d' % i)
109110
b1_list.append(b1)
110111
b2_list.append(b2)
111112
net = tl.layers.merge_networks([b1_list[-1], b2_list[-1]])

‎setup.cfg

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
[yapf]
2+
based_on_style=google
3+
4+
# The number of columns to use for indentation.
5+
indent_width = 4
6+
7+
# The column limit.
8+
column_limit=120
9+
10+
# Place each dictionary entry onto its own line.
11+
each_dict_entry_on_separate_line = True

‎train.py

Lines changed: 35 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,15 @@
1414
from config import config
1515
from models import model
1616
from pycocotools.coco import maskUtils
17-
from tensorlayer.prepro import (
18-
keypoint_random_crop, keypoint_random_flip, keypoint_random_resize,
19-
keypoint_random_resize_shortestedge, keypoint_random_rotate)
20-
from utils import (PoseInfo, draw_intermedia_results, get_heatmap,
21-
get_vectormap, load_mscoco_dataset)
17+
from tensorlayer.prepro import (keypoint_random_crop, keypoint_random_flip, keypoint_random_resize,
18+
keypoint_random_resize_shortestedge, keypoint_random_rotate)
19+
from utils import (PoseInfo, draw_intermedia_results, get_heatmap, get_vectormap, load_mscoco_dataset)
2220

2321
tf.logging.set_verbosity(tf.logging.DEBUG)
2422
tl.logging.set_verbosity(tl.logging.DEBUG)
2523

26-
tl.files.exists_or_mkdir(
27-
config.LOG.vis_path, verbose=False) # to save visualization results
28-
tl.files.exists_or_mkdir(
29-
config.MODEL.model_path, verbose=False) # to save model files
24+
tl.files.exists_or_mkdir(config.LOG.vis_path, verbose=False) # to save visualization results
25+
tl.files.exists_or_mkdir(config.MODEL.model_path, verbose=False) # to save model files
3026

3127
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
3228
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
@@ -83,16 +79,12 @@ def _data_aug_fn(image, ground_truth):
8379
mask_miss = np.bitwise_and(mask_miss, bin_mask)
8480

8581
# image data augmentation
86-
image, annos, mask_miss = keypoint_random_resize(
87-
image, annos, mask_miss, zoom_range=(0.8, 1.2))
88-
image, annos, mask_miss = keypoint_random_rotate(
89-
image, annos, mask_miss, rg=15.0)
90-
image, annos, mask_miss = keypoint_random_flip(
91-
image, annos, mask_miss, prob=0.5)
82+
image, annos, mask_miss = keypoint_random_resize(image, annos, mask_miss, zoom_range=(0.8, 1.2))
83+
image, annos, mask_miss = keypoint_random_rotate(image, annos, mask_miss, rg=15.0)
84+
image, annos, mask_miss = keypoint_random_flip(image, annos, mask_miss, prob=0.5)
9285
image, annos, mask_miss = keypoint_random_resize_shortestedge(
9386
image, annos, mask_miss, min_size=(hin, win)) # TODO: give size
94-
image, annos, mask_miss = keypoint_random_crop(
95-
image, annos, mask_miss, size=(hin, win)) # TODO: give size
87+
image, annos, mask_miss = keypoint_random_crop(image, annos, mask_miss, size=(hin, win)) # TODO: give size
9688

9789
# generate result maps including keypoints heatmap, pafs and mask
9890
h, w, _ = np.shape(image)
@@ -107,8 +99,7 @@ def _data_aug_fn(image, ground_truth):
10799
image = image * np.repeat(img_mask, 3, 2)
108100

109101
resultmap = np.array(resultmap, dtype=np.float32)
110-
mask_miss = cv2.resize(
111-
mask_miss, (hout, wout), interpolation=cv2.INTER_AREA)
102+
mask_miss = cv2.resize(mask_miss, (hout, wout), interpolation=cv2.INTER_AREA)
112103
mask_miss = np.array(mask_miss, dtype=np.float32)
113104
return image, resultmap, mask_miss
114105

@@ -118,8 +109,7 @@ def _map_fn(img_list, annos):
118109
image = tf.read_file(img_list)
119110
image = tf.image.decode_jpeg(image, channels=3) # get RGB with 0~1
120111
image = tf.image.convert_image_dtype(image, dtype=tf.float32)
121-
image, resultmap, mask = tf.py_func(_data_aug_fn, [image, annos],
122-
[tf.float32, tf.float32, tf.float32])
112+
image, resultmap, mask = tf.py_func(_data_aug_fn, [image, annos], [tf.float32, tf.float32, tf.float32])
123113
return image, resultmap, mask
124114

125115

@@ -136,8 +126,7 @@ def _map_fn(img_list, annos):
136126
train_mask_list = train_data.get_mask()
137127
# train_targets = list(zip(train_objs_info_list, train_mask_list))
138128
if len(train_imgs_file_list) != len(train_objs_info_list):
139-
raise Exception(
140-
"number of training images and annotations do not match")
129+
raise Exception("number of training images and annotations do not match")
141130
else:
142131
print("number of training images {}".format(len(train_imgs_file_list)))
143132

@@ -160,21 +149,17 @@ def _map_fn(img_list, annos):
160149
your_objs_info_list = your_data.get_joint_list()
161150
your_mask_list = your_data.get_mask()
162151
if len(your_imgs_file_list) != len(your_objs_info_list):
163-
raise Exception(
164-
"number of customized images and annotations do not match")
152+
raise Exception("number of customized images and annotations do not match")
165153
else:
166-
print("number of customized images {}".format(
167-
len(your_imgs_file_list)))
154+
print("number of customized images {}".format(len(your_imgs_file_list)))
168155

169156
# choice dataset for training
170157
# 1. only coco training set
171158
# imgs_file_list = train_imgs_file_list
172159
# train_targets = list(zip(train_objs_info_list, train_mask_list))
173160
# 2. your customized data from "data/your_data" and coco training set
174161
imgs_file_list = train_imgs_file_list + your_imgs_file_list
175-
train_targets = list(
176-
zip(train_objs_info_list + your_objs_info_list,
177-
train_mask_list + your_mask_list))
162+
train_targets = list(zip(train_objs_info_list + your_objs_info_list, train_mask_list + your_mask_list))
178163

179164
# define data augmentation
180165
def generator():
@@ -183,8 +168,7 @@ def generator():
183168
for _input, _target in zip(imgs_file_list, train_targets):
184169
yield _input.encode('utf-8'), cPickle.dumps(_target)
185170

186-
dataset = tf.data.Dataset().from_generator(
187-
generator, output_types=(tf.string, tf.string))
171+
dataset = tf.data.Dataset().from_generator(generator, output_types=(tf.string, tf.string))
188172
dataset = dataset.map(_map_fn, num_parallel_calls=8)
189173
dataset = dataset.shuffle(buffer_size=2046)
190174
dataset = dataset.repeat(n_epoch)
@@ -197,19 +181,14 @@ def generator():
197181
# Train with placeholder can help your to check the data easily.
198182
# define model architecture
199183
x = tf.placeholder(tf.float32, [None, hin, win, 3], "image")
200-
confs = tf.placeholder(tf.float32, [None, hout, wout, n_pos],
201-
"confidence_maps")
202-
pafs = tf.placeholder(tf.float32, [None, hout, wout, n_pos * 2],
203-
"pafs")
184+
confs = tf.placeholder(tf.float32, [None, hout, wout, n_pos], "confidence_maps")
185+
pafs = tf.placeholder(tf.float32, [None, hout, wout, n_pos * 2], "pafs")
204186
# if the people does not have keypoints annotations, ignore the area
205-
img_mask1 = tf.placeholder(tf.float32, [None, hout, wout, n_pos],
206-
'img_mask1')
207-
img_mask2 = tf.placeholder(tf.float32, [None, hout, wout, n_pos * 2],
208-
'img_mask2')
187+
img_mask1 = tf.placeholder(tf.float32, [None, hout, wout, n_pos], 'img_mask1')
188+
img_mask2 = tf.placeholder(tf.float32, [None, hout, wout, n_pos * 2], 'img_mask2')
209189
num_images = np.shape(imgs_file_list)[0]
210190

211-
cnn, b1_list, b2_list, net = model(x, n_pos, img_mask1, img_mask2,
212-
True, False)
191+
cnn, b1_list, b2_list, net = model(x, n_pos, img_mask1, img_mask2, True, False)
213192

214193
# define loss
215194
losses = []
@@ -218,12 +197,8 @@ def generator():
218197
stage_losses = []
219198
L2 = 0.0
220199
for idx, (l1, l2) in enumerate(zip(b1_list, b2_list)):
221-
loss_l1 = tf.nn.l2_loss(
222-
(tf.concat(l1.outputs, axis=0) - tf.concat(confs, axis=0)) *
223-
img_mask1)
224-
loss_l2 = tf.nn.l2_loss(
225-
(tf.concat(l2.outputs, axis=0) - tf.concat(pafs, axis=0)) *
226-
img_mask2)
200+
loss_l1 = tf.nn.l2_loss((tf.concat(l1.outputs, axis=0) - tf.concat(confs, axis=0)) * img_mask1)
201+
loss_l2 = tf.nn.l2_loss((tf.concat(l2.outputs, axis=0) - tf.concat(pafs, axis=0)) * img_mask2)
227202
losses.append(tf.reduce_mean([loss_l1, loss_l2]))
228203
stage_losses.append(loss_l1 / batch_size)
229204
stage_losses.append(loss_l2 / batch_size)
@@ -237,15 +212,14 @@ def generator():
237212
total_loss = tf.reduce_sum(losses) / batch_size + L2
238213

239214
global_step = tf.Variable(1, trainable=False)
240-
print('Config:', 'n_epoch: ', n_epoch, 'batch_size: ', batch_size,
241-
'base_lr: ', base_lr, 'step_size: ', step_size)
215+
print('Config:', 'n_epoch: ', n_epoch, 'batch_size: ', batch_size, 'base_lr: ', base_lr, 'step_size: ',
216+
step_size)
242217
with tf.variable_scope('learning_rate'):
243218
lr_v = tf.Variable(base_lr, trainable=False)
244219

245220
opt = tf.train.MomentumOptimizer(lr_v, 0.9)
246221
train_op = opt.minimize(total_loss, global_step=global_step)
247-
config = tf.ConfigProto(
248-
allow_soft_placement=True, log_device_placement=False)
222+
config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)
249223

250224
# start training
251225
with tf.Session(config=config) as sess:
@@ -292,14 +266,8 @@ def generator():
292266
# os.path.join(config.LOG.vis_path, 'data_aug_{}.png'.format(i))
293267
# tl.file.save_image()
294268

295-
[
296-
_, the_loss, loss_ll, L2_reg, conf_result, weight_norm,
297-
paf_result
298-
] = sess.run(
299-
[
300-
train_op, total_loss, stage_losses, L2, last_conf, L2,
301-
last_paf
302-
],
269+
[_, the_loss, loss_ll, L2_reg, conf_result, weight_norm, paf_result] = sess.run(
270+
[train_op, total_loss, stage_losses, L2, last_conf, L2, last_paf],
303271
feed_dict={
304272
x: x_,
305273
confs: confs_,
@@ -308,35 +276,25 @@ def generator():
308276
img_mask2: mask2
309277
})
310278

311-
tstring = time.strftime('%d-%m %H:%M:%S',
312-
time.localtime(time.time()))
279+
tstring = time.strftime('%d-%m %H:%M:%S', time.localtime(time.time()))
313280
lr = sess.run(lr_v)
314-
print(
315-
'Total Loss at iteration {} is: {} Learning rate {:10e} weight_norm {:10e} Time: {}'.
316-
format(gs_num, the_loss, lr, weight_norm, tstring))
281+
print('Total Loss at iteration {} is: {} Learning rate {:10e} weight_norm {:10e} Time: {}'.format(
282+
gs_num, the_loss, lr, weight_norm, tstring))
317283
for ix, ll in enumerate(loss_ll):
318-
print('Network#', ix, 'For Branch', ix % 2 + 1, 'Loss:',
319-
ll)
284+
print('Network#', ix, 'For Branch', ix % 2 + 1, 'Loss:', ll)
320285

321286
# save some intermedian results
322287
if (gs_num != 0) and (gs_num % 1 == 0): # save_interval == 0):
323-
draw_intermedia_results(x_, confs_, conf_result, pafs_,
324-
paf_result, mask, 'train')
288+
draw_intermedia_results(x_, confs_, conf_result, pafs_, paf_result, mask, 'train')
325289
# np.save(config.LOG.vis_path + 'image' + str(gs_num) + '.npy', x_)
326290
# np.save(config.LOG.vis_path + 'heat_ground' + str(gs_num) + '.npy', confs_)
327291
# np.save(config.LOG.vis_path + 'heat_result' + str(gs_num) + '.npy', conf_result)
328292
# np.save(config.LOG.vis_path + 'paf_ground' + str(gs_num) + '.npy', pafs_)
329293
# np.save(config.LOG.vis_path + 'mask' + str(gs_num) + '.npy', mask)
330294
# np.save(config.LOG.vis_path + 'paf_result' + str(gs_num) + '.npy', paf_result)
331295
tl.files.save_npz_dict(
332-
net.all_params,
333-
os.path.join(model_path,
334-
'pose' + str(gs_num) + '.npz'),
335-
sess=sess)
336-
tl.files.save_npz_dict(
337-
net.all_params,
338-
os.path.join(model_path, 'pose.npz'),
339-
sess=sess)
296+
net.all_params, os.path.join(model_path, 'pose' + str(gs_num) + '.npz'), sess=sess)
297+
tl.files.save_npz_dict(net.all_params, os.path.join(model_path, 'pose.npz'), sess=sess)
340298
if gs_num > 3000001:
341299
break
342300
elif config.TRAIN.train_mode == 'dataset': # TODO

0 commit comments

Comments
 (0)