def main():
use_gpu = False
if use_gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
else:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
input_img = tf.placeholder(tf.float32, shape=[1, None, None, 3])
# _1, _2, cpm, paf = light_openpose(input_img, is_training=False)
cpm, paf = lightweight_openpose(input_img, num_pafs=26, num_joints=14, is_training=False)
saver = tf.train.Saver()
total_img = 0
total_time = 0
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
saver.restore(sess, params['test_model'])
print('#---------Successfully loaded trained model.---------#')
save(sess)
if params['img_path'] is not None:
img_ori = cv2.imread(params['img_path'])
print(np.shape(img_ori))
img_data = cv2.cvtColor(img_ori, code=cv2.COLOR_BGR2RGB)
img_data = cv2.resize(img_data, (256, 256))
img = img_data / 255.
start_time = time.time()
heatmap, _paf = sess.run([cpm, paf], feed_dict={input_img: [img]})
end_time = time.time()
print(heatmap.shape)
print(_paf.shape)
canvas, joint_list, person_to_joint_assoc, joints = decode_pose(img_data, params, heatmap[0], _paf[0])
decode_time = time.time()
print('inference + decode time == {}'.format(decode_time - start_time))
total_img += 1
total_time += (end_time - start_time)
canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
cv2_imshow(canvas)
else:
print('Nothing to process.')
def main():
use_gpu = False
if use_gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
else:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
input_img = tf.placeholder(tf.float32, shape=[1, None, None, 3])
# _1, _2, cpm, paf = light_openpose(input_img, is_training=False)
cpm, paf = lightweight_openpose(input_img, num_pafs=26, num_joints=14, is_training=False)
saver = tf.train.Saver()
total_img = 0
total_time = 0
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
saver.restore(sess, params['test_model'])
print('#---------Successfully loaded trained model.---------#')
if 'video_path'in params.keys() and params['video_path'] is not None:
# video_capture = cv2.VideoCapture('rtsp://*****:*****@10.24.1.238')
video_capture = cv2.VideoCapture(params['video_path'])
fps = video_capture.get(cv2.CAP_PROP_FPS)
start_second = 0
start_frame = fps * start_second
video_capture.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
while True:
retval, img_data = video_capture.read()
if not retval:
break
img_data = cv2.cvtColor(img_data, code=cv2.COLOR_BGR2RGB)
orih, oriw, c = img_data.shape
img = cv2.resize(img_data, (256, 256)) / 255.
start_time = time.time()
heatmap, _paf = sess.run([cpm, paf], feed_dict={input_img: [img]})
end_time = time.time()
canvas, joint_list, person_to_joint_assoc = decode_pose(img, params, heatmap[0], _paf[0])
decode_time = time.time()
print ('inference + decode time == {}'.format(decode_time - start_time))
total_img += 1
total_time += (end_time-start_time)
canvas = canvas.astype(np.float32)
canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
cv2.imshow('result', canvas)
cv2.waitKey(1)
elif params['img_path'] is not None:
for img_name in os.listdir(params['img_path']):
if img_name.split('.')[-1] != 'jpg':
continue
img_data = cv2.imread(os.path.join(params['img_path'], img_name))
img_data = cv2.cvtColor(img_data, code=cv2.COLOR_BGR2RGB)
img = cv2.resize(img_data, (256, 256)) / 255.
start_time = time.time()
heatmap, _paf = sess.run([cpm, paf], feed_dict={input_img: [img]})
end_time = time.time()
canvas, joint_list, person_to_joint_assoc = decode_pose(img, params, heatmap[0], _paf[0])
canvas = canvas.astype(np.float32)
canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
decode_time = time.time()
print('inference + decode time == {}'.format(decode_time - start_time))
total_img += 1
total_time += (end_time - start_time)
cv2.imshow('result', canvas)
cv2.waitKey(0)
else:
print('Nothing to process.')
def model_fn(features, labels, mode, params):
# get model output
features = tf.reshape(features, [-1, params['height'],params['width'], 3])
gt_cpms = labels[..., :params['num_kps']]
gt_pafs = labels[..., params['num_kps']:params['num_kps'] + params['paf']]
mask = labels[..., params['num_kps'] + params['paf']:]
mask = tf.reshape(mask, [-1, params['height']//params['scale'], params['width']//params['scale'], 1])
cpm, paf = lightweight_openpose(inputs=features, num_joints=params['num_kps'], num_pafs=params['paf'], is_training=True)
predictions = {
'pred_heatmap': cpm,
'pred_paf': paf
}
tf.summary.image('img', features, max_outputs=3)
tf.summary.image('pred_hmap', tf.reduce_sum(cpm, axis=3, keepdims=True), max_outputs=3)
tf.summary.image('gt_hmap', tf.reduce_sum(gt_cpms, axis=3, keepdims=True), max_outputs=3)
tf.summary.image('gt_paf', tf.expand_dims(
(gt_pafs[..., 0] - tf.reduce_min(gt_pafs[..., 0])) / (tf.reduce_max(gt_pafs[..., 0]) - tf.reduce_min(gt_pafs[..., 0])),
axis=3
), max_outputs=3)
tf.summary.image('pred_paf', tf.expand_dims(
(paf[..., 0] - tf.reduce_min(paf[..., 0])) / (tf.reduce_max(paf[..., 0]) - tf.reduce_min(paf[..., 0])),
axis=3
), max_outputs=3)
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
export_outputs={
'predict': tf.estimator.export.PredictOutput(predictions)
})
cpm_mask = tf.concat([mask for i in range(params['num_kps'])], axis=-1)
paf_mask = tf.concat([mask for i in range(params['paf'])], axis=-1)
cpm = tf.where(cpm_mask > 0, cpm, cpm * 0)
paf = tf.where(paf_mask > 0, paf, paf * 0)
gt_cpms = tf.where(cpm_mask > 0, gt_cpms, gt_cpms * 0)
gt_pafs = tf.where(paf_mask > 0, gt_pafs, gt_pafs * 0)
loss = tf.nn.l2_loss(cpm - gt_cpms) + tf.nn.l2_loss(paf - gt_pafs) * 2
tf.identity(loss, name='loss')
tf.summary.scalar('loss', loss)
if mode == tf.estimator.ModeKeys.EVAL:
metrics_dict = {
'heatmap': tf.metrics.mean_squared_error(labels=gt_cpms, predictions=predictions['pred_heatmap']),
'paf': tf.metrics.mean_squared_error(labels=gt_pafs, predictions=predictions['pred_paf'])
}
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
eval_metric_ops=metrics_dict
)
if mode == tf.estimator.ModeKeys.TRAIN:
global_step = tf.train.get_or_create_global_step()
# step lr
# values = [params['lr'], 0.1*params['lr'], 0.01*params['lr'], 0.001*params['lr']]
# boundaries = [params['train_nums']*50, params['train_nums']*100, params['train_nums']*150]
# learning_rate = tf.train.piecewise_constant(global_step, boundaries, values)
# constant lr
learning_rate = tf.Variable(params['lr'], trainable=False, name='lr')
tf.identity(learning_rate, name='lr')
tf.summary.scalar('lr', learning_rate)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate, epsilon=1e-5)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(loss, global_step=global_step)
else:
train_op = None
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op
)
def main():
use_gpu = params['gpu']
if use_gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
else:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
with open(params['json_file'], encoding='utf-8') as fp:
labels = json.load(fp)
input_img = tf.placeholder(tf.float32, shape=[1, None, None, 3])
cpm, paf = lightweight_openpose(input_img,
num_pafs=26,
num_joints=14,
is_training=False)
saver = tf.train.Saver()
total_img = 0
predictions = []
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
saver.restore(sess, params['test_model'])
print('#---------Successfully loaded trained model.---------#')
for label in labels:
img_id = label['image_id']
img_ori = cv2.imread(
os.path.join(params['img_path'], img_id + '.jpg'))
img_data = cv2.cvtColor(img_ori, code=cv2.COLOR_BGR2RGB)
img_data = cv2.resize(img_data,
(params['input_size'], params['input_size']))
img = img_data / 255.
heatmap, _paf = sess.run([cpm, paf], feed_dict={input_img: [img]})
canvas, joint_list, person_to_joint_assoc, joints = decode_pose(
img_data, params, heatmap[0], _paf[0])
predict = label
predict['image_id'] = img_id
kps = {}
human = 1
factorY = img_ori.shape[0] / img_data.shape[0]
factorX = img_ori.shape[1] / img_data.shape[1]
for joint in joints:
for i in range(14):
joint[3 * i] *= factorX
joint[3 * i + 1] *= factorY
if type(joint) == type([]):
kps['human' + str(human)] = joint
else:
kps['human' + str(human)] = joint.tolist()
human += 1
# print (human)
# print (kps)
# for person, joint in kps.items():
# assert len(joint) == 14 * 3
# for i in range(14):
# x = int(joint[i*3])
# y = int(joint[i*3+1])
# cv2.circle(img_ori, (x, y), 3, (255, 255, 255), thickness=-1)
# cv2.putText(img_ori, str(i), (x, y),cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 1)
# cv2.imshow('test', img_ori)
# cv2.waitKey(0)
predict['keypoint_annotations'] = kps
predictions.append(predict)
# break
total_img += 1
print('processing number: {}'.format(total_img))
# break
with open(params['save_json'], 'w') as fw:
json.dump(predictions, fw)