Towards TF2.0

Author: rigolepe

docker_run.sh

@@ -25,4 +25,4 @@ else
   gpu_opts=""
 fi
 
-docker run $gpu_opts -it --rm -v $WORK:/work "$image" python "$@"
+docker run $gpu_opts -it --rm -v $WORK:/work "$image" "$@"

get_test_images_run.sh

@@ -1,3 +1,3 @@
 #!/usr/bin/env bash
 
-./docker_run.sh get_test_images.py
+./docker_run.sh python get_test_images.py

image_demo_run.sh

@@ -1,3 +1,3 @@
 #!/usr/bin/env bash
 
-./docker_run.sh image_demo.py --model 101 --image_dir ./images --output_dir ./output
+./docker_run.sh python image_demo.py --model 101 --image_dir ./images --output_dir ./output

image_test.py

@@ -39,7 +39,7 @@ def main():
         img = cv2.imread(f)
         pose_scores, keypoint_scores, keypoint_coords = posenet.estimate_multiple_poses(img)
         img_poses = posenet.draw_poses(img, pose_scores, keypoint_scores, keypoint_coords)
-        posenet.print_scores(img, pose_scores, keypoint_scores, keypoint_coords)
+        posenet.print_scores(f, pose_scores, keypoint_scores, keypoint_coords)
         cv2.imwrite(os.path.join(args.output_dir, os.path.relpath(f, args.image_dir)), img_poses)
 
     print('Average FPS:', len(filenames) / (time.time() - start))

image_test_run.sh

@@ -1,3 +1,3 @@
 #!/usr/bin/env bash
 
-./docker_run.sh image_test.py --model 101 --image_dir ./images --output_dir ./output
+./docker_run.sh python image_test.py --model 101 --image_dir ./images --output_dir ./output

inspect_saved_model.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+FOLDER=$1
+
+# e.g.: $> ./inspect_saved_model.sh _tf_models/posenet/mobilenet_v1_100/stride16
+./docker_run.sh saved_model_cli show --dir "$FOLDER" --all

posenet/base_model.py

@@ -4,16 +4,19 @@
 
 class BaseModel(ABC):
 
-    def __init__(self, sess, input_tensor_name, output_tensor_names, output_stride):
+    # keys for the output_tensor_names map
+    HEATMAP_KEY = "heatmap"
+    OFFSETS_KEY = "offsets"
+    DISPLACEMENT_FWD_KEY = "displacement_fwd"
+    DISPLACEMENT_BWD_KEY = "displacement_bwd"
+
+    def __init__(self, model_function, output_tensor_names, output_stride):
         self.output_stride = output_stride
-        self.sess = sess
-        self.input_tensor_name = input_tensor_name
-        self.output_tensors = [
-            tf.sigmoid(sess.graph.get_tensor_by_name(output_tensor_names['heatmap']), 'heatmap'),  # sigmoid!!!
-            sess.graph.get_tensor_by_name(output_tensor_names['offsets']),
-            sess.graph.get_tensor_by_name(output_tensor_names['displacement_fwd']),
-            sess.graph.get_tensor_by_name(output_tensor_names['displacement_bwd'])
-        ]
+        self.output_tensor_names = output_tensor_names
+        self.model_function = model_function
+        # self.sess = sess
+        # self.input_tensor_name = input_tensor_name
+        # self.output_tensors = output_tensors
 
     def valid_resolution(self, width, height):
         # calculate closest smaller width and height that is divisible by the stride after subtracting 1 (for the bias?)
@@ -27,11 +30,24 @@ def preprocess_input(self, image):
 
     def predict(self, image):
         input_image, image_scale = self.preprocess_input(image)
-        heatmap_result, offsets_result, displacement_fwd_result, displacement_bwd_result = self.sess.run(
-            self.output_tensors,
-            feed_dict={self.input_tensor_name: input_image}
-        )
-        return heatmap_result, offsets_result, displacement_fwd_result, displacement_bwd_result, image_scale
+
+        input_image = tf.convert_to_tensor(input_image, dtype=tf.float32)
+
+        result = self.model_function(input_image)
+
+        heatmap_result = result[self.output_tensor_names[self.HEATMAP_KEY]]
+        offsets_result = result[self.output_tensor_names[self.OFFSETS_KEY]]
+        displacement_fwd_result = result[self.output_tensor_names[self.DISPLACEMENT_FWD_KEY]]
+        displacement_bwd_result  = result[self.output_tensor_names[self.DISPLACEMENT_BWD_KEY]]
+
+
+            # self.sess.run(
+            # self.output_tensors,
+            # feed_dict={self.input_tensor_name: input_image}
+        # )
+
+        return tf.sigmoid(heatmap_result), offsets_result, displacement_fwd_result, displacement_bwd_result, image_scale
 
     def close(self):
-        self.sess.close()
+        # self.sess.close()
+        return

posenet/converter/config.py

@@ -6,5 +6,5 @@
 
 def load_config(config_name='config.yaml'):
     cfg_f = open(os.path.join(BASE_DIR, config_name), "r+")
-    cfg = yaml.load(cfg_f)
+    cfg = yaml.load(cfg_f, Loader=yaml.FullLoader)
     return cfg

posenet/converter/tfjs2tf.py

@@ -4,13 +4,46 @@
 import posenet.converter.tfjsdownload as tfjsdownload
 
 
+def __tensor_info_def(sess, tensor_names):
+    signatures = {}
+    for tensor_name in tensor_names:
+        tensor = sess.graph.get_tensor_by_name(tensor_name)
+        tensor_info = tf.compat.v1.saved_model.build_tensor_info(tensor)
+        signatures[tensor_name] = tensor_info
+    return signatures
+
+
 def convert(model, neuralnet, model_variant):
     model_cfg = tfjsdownload.model_config(model, neuralnet, model_variant)
     model_file_path = os.path.join(model_cfg['tfjs_dir'], model_cfg['filename'])
     if not os.path.exists(model_file_path):
         print('Cannot find tfjs model path %s, downloading tfjs model...' % model_file_path)
         tfjsdownload.download_tfjs_model(model, neuralnet, model_variant)
-    tfjs.api.graph_model_to_saved_model(model_cfg['tfjs_dir'], model_cfg['tf_dir'], ['serve'])
+
+    # 'graph_model_to_saved_model' doesn't store the signature for the model!
+    # tfjs.api.graph_model_to_saved_model(model_cfg['tfjs_dir'], model_cfg['tf_dir'], ['serve'])
+    # so we do it manually here:
+
+    # see: https://www.programcreek.com/python/example/104885/tensorflow.python.saved_model.signature_def_utils.build_signature_def
+    graph = tfjs.api.load_graph_model(model_cfg['tfjs_dir'])
+    builder = tf.compat.v1.saved_model.Builder(model_cfg['tf_dir'])
+
+    with tf.compat.v1.Session(graph=graph) as sess:
+        input_tensor_names = tfjs.util.get_input_tensors(graph)
+        output_tensor_names = tfjs.util.get_output_tensors(graph)
+
+        signature_inputs = __tensor_info_def(sess, input_tensor_names)
+        signature_outputs = __tensor_info_def(sess, output_tensor_names)
+
+        method_name = tf.compat.v1.saved_model.signature_constants.PREDICT_METHOD_NAME
+        signature_def = tf.compat.v1.saved_model.build_signature_def(inputs=signature_inputs,
+                                                                     outputs=signature_outputs,
+                                                                     method_name=method_name)
+        signature_map = {tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def}
+        builder.add_meta_graph_and_variables(sess=sess,
+                                             tags=['serve'],
+                                             signature_def_map=signature_map)
+    return builder.save()
 
 
 def list_tensors(model, neuralnet, model_variant):

posenet/mobilenet.py

@@ -5,8 +5,8 @@
 
 class MobileNet(BaseModel):
 
-    def __init__(self, sess, input_tensor_name, output_tensor_names, output_stride):
-        super().__init__(sess, input_tensor_name, output_tensor_names, output_stride)
+    def __init__(self, model_function, output_tensor_names, output_stride):
+        super().__init__(model_function, output_tensor_names, output_stride)
 
     def preprocess_input(self, image):
         target_width, target_height = self.valid_resolution(image.shape[1], image.shape[0])