import urllib.request
import posixpath
import json
import os
from posenet.converter.config import load_config
CFG = load_config()
GOOGLE_CLOUD_STORAGE_DIR = CFG['GOOGLE_CLOUD_STORAGE_DIR']
CHECKPOINTS = CFG['checkpoints']
CHK = CFG['chk']
def download_file(checkpoint, filename, base_dir):
url = posixpath.join(GOOGLE_CLOUD_STORAGE_DIR, checkpoint, filename)
urllib.request.urlretrieve(url, os.path.join(base_dir, checkpoint,
filename))
def download(checkpoint, base_dir='./weights/'):
save_dir = os.path.join(base_dir, checkpoint)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
download_file(checkpoint, 'manifest.json', base_dir)
f = open(os.path.join(save_dir, 'manifest.json'), 'r')
json_dict = json.load(f)
for x in json_dict:
filename = json_dict[x]['filename']
def convert(model_id, model_dir, check=False):
cfg = load_config()
checkpoints = cfg['checkpoints']
image_size = cfg['imageSize']
output_stride = cfg['outputStride']
chkpoint = checkpoints[model_id]
if chkpoint == 'mobilenet_v1_050':
mobile_net_arch = cfg['mobileNet50Architecture']
elif chkpoint == 'mobilenet_v1_075':
mobile_net_arch = cfg['mobileNet75Architecture']
else:
mobile_net_arch = cfg['mobileNet100Architecture']
width = image_size
height = image_size
if not os.path.exists(model_dir):
os.makedirs(model_dir)
cg = tf.Graph()
with cg.as_default():
layers = to_output_strided_layers(mobile_net_arch, output_stride)
variables = load_variables(chkpoint)
init = tf.global_variables_initializer()
with tf.Session() as sess:
sess.run(init)
saver = tf.train.Saver()
image_ph = tf.placeholder(tf.float32,
shape=[1, None, None, 3],
name='image')
outputs = build_network(image_ph, layers, variables)
sess.run(
[outputs],
feed_dict={
image_ph:
[np.ndarray(shape=(height, width, 3), dtype=np.float32)]
})
save_path = os.path.join(model_dir, 'checkpoints',
'model-%s.ckpt' % chkpoint)
if not os.path.exists(os.path.dirname(save_path)):
os.makedirs(os.path.dirname(save_path))
checkpoint_path = saver.save(sess, save_path, write_state=False)
tf.train.write_graph(cg, model_dir, "model-%s.pbtxt" % chkpoint)
output_node_names = 'heatmap,offset_2,displacement_fwd_2,displacement_bwd_2'
# Freeze graph and write our final model file
freeze_graph(input_graph=os.path.join(model_dir,
"model-%s.pbtxt" % chkpoint),
input_saver="",
input_binary=False,
input_checkpoint=checkpoint_path,
output_node_names=output_node_names,
restore_op_name="save/restore_all",
filename_tensor_name="save/Const:0",
output_graph=os.path.join(model_dir,
"model-%s.pb" % chkpoint),
clear_devices=True,
initializer_nodes="")
if check and os.path.exists("./images/tennis_in_crowd.jpg"):
# Result TF
input_image = _read_imgfile("./images/tennis_in_crowd.jpg",
width, height)
input_image = np.array(input_image, dtype=np.float32)
input_image = input_image.reshape(1, height, width, 3)
heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = sess.run(
outputs, feed_dict={image_ph: input_image})
print("Test image stats")
print(input_image)
print(input_image.shape)
print(np.mean(input_image))
heatmaps_result = heatmaps_result[0]
print("Heatmaps")
print(heatmaps_result[0:1, 0:1, :])
print(heatmaps_result.shape)
print(np.mean(heatmaps_result))
#Result CoreML
out = coreml_model.predict({'image__0': img})['heatmap__0']
print("#output coreml result.")
print(out.shape)
print(np.transpose(out))
print(out)
# print(out[:, 0:1, 0:1])
print(np.mean(out))
def convert(model_id, model_dir, check=False):
cfg = load_config()
checkpoints = cfg['checkpoints']
image_size = cfg['imageSize']
output_stride = cfg['outputStride']
chkpoint = checkpoints[model_id]
if chkpoint == 'mobilenet_v1_050':
mobile_net_arch = cfg['mobileNet50Architecture']
elif chkpoint == 'mobilenet_v1_075':
mobile_net_arch = cfg['mobileNet75Architecture']
else:
mobile_net_arch = cfg['mobileNet100Architecture']
width = image_size
height = image_size
if not os.path.exists(model_dir):
os.makedirs(model_dir)
cg = tf.Graph()
with cg.as_default():
layers = to_output_strided_layers(mobile_net_arch, output_stride)
variables = load_variables(chkpoint)
init = tf.compat.v1.global_variables_initializer()
with tf.compat.v1.Session() as sess:
sess.run(init)
saver = tf.compat.v1.train.Saver()
image_ph = tf.compat.v1.placeholder(tf.float32,
shape=[1, None, None, 3],
name='image')
outputs = build_network(image_ph, layers, variables)
sess.run(
[outputs],
feed_dict={
image_ph:
[np.ndarray(shape=(height, width, 3), dtype=np.float32)]
})
save_path = os.path.join(model_dir, 'checkpoints',
'model-%s.ckpt' % chkpoint)
if not os.path.exists(os.path.dirname(save_path)):
os.makedirs(os.path.dirname(save_path))
checkpoint_path = saver.save(sess, save_path, write_state=False)
tf.io.write_graph(cg, model_dir, "model-%s.pbtxt" % chkpoint)
# Freeze graph and write our final model file
frozen_graph = freeze_session(sess, None, [
'heatmap', 'offset_2', 'displacement_fwd_2',
'displacement_bwd_2'
], True)
tf.compat.v1.train.write_graph(frozen_graph,
'./',
os.path.join(
model_dir,
"model-%s.pb" % chkpoint),
as_text=False)
if os.path.exists("./images/tennis_in_crowd.jpg"):
# Result
input_image = _read_imgfile("./images/tennis_in_crowd.jpg",
width, height)
input_image = np.array(input_image, dtype=np.float32)
input_image = input_image.reshape(1, height, width, 3)
heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = sess.run(
outputs, feed_dict={image_ph: input_image})
print("Test image stats")
print(input_image)
print(input_image.shape)
print(np.mean(input_image))
heatmaps_result = heatmaps_result[0]
print("Heatmaps")
print(heatmaps_result[0:1, 0:1, :])
print(heatmaps_result.shape)
print(np.mean(heatmaps_result))
