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))