Refactoring everything to the new saved model.

rigolepe · rigolepe · commit 7a2f5fc9add2 · 2020-01-03T00:00:02.000+01:00
diff --git a/.gitignore b/.gitignore
@@ -1,8 +1,11 @@
 images/*
 output/*
+output_old/*
 .idea/*
 .idea
 _models/*
+_tf_models/*
+_tfjs_models/*
 _posenet_weights/*
 env.sh
 
diff --git a/Dockerfile b/Dockerfile
@@ -1,5 +1,7 @@
-FROM tensorflow/tensorflow:nightly-py3-jupyter
-# This is the CPU version!
+# default image version, override using --build-arg IMAGE_VERSION=otherversion
+ARG IMAGE_VERSION=nightly-py3-jupyter
+FROM tensorflow/tensorflow:$IMAGE_VERSION
+# The default version is the CPU version!
 # see: https://www.tensorflow.org/install/docker
 # see: https://hub.docker.com/r/tensorflow/tensorflow/
 
diff --git a/Dockerfile-gpu b/Dockerfile-gpu
diff --git a/benchmark.py b/benchmark.py
@@ -16,8 +16,7 @@
 def main():
 
     with tf.compat.v1.Session() as sess:
-        model_cfg, model_outputs = posenet.load_model(args.model, sess)
-        output_stride = model_cfg['output_stride']
+        output_stride, model_outputs = posenet.load_model(args.model, sess)
         num_images = args.num_images
 
         filenames = [
diff --git a/docker_img_build.sh b/docker_img_build.sh
@@ -8,12 +8,12 @@ fi
 
 if [ "$1" = "GPU" ]; then
   image="posenet-python-gpu"
-  dockerfile="Dockerfile-gpu"
+  version="--build-arg IMAGE_VERSION=2.0.0-gpu-py3-jupyter"
 else
   image="posenet-python-cpu"
-  dockerfile="Dockerfile"
+  version=""
 fi
 
 docker rmi -f "$image"
 
-docker build -t "$image" -f "$dockerfile" .
+docker build -t "$image" $version .
diff --git a/docker_run.sh b/docker_run.sh
@@ -20,8 +20,10 @@ echo "device is: $POSENET_PYTHON_DEVICE"
 
 if [ "$POSENET_PYTHON_DEVICE" = "GPU" ]; then
   image="posenet-python-gpu"
+  gpu_opts="--gpus all"
 else
   image="posenet-python-cpu"
+  gpu_opts=""
 fi
 
-docker run --gpus all -it --rm -v $WORK:/work "$image" python "$@"
+docker run $gpu_opts -it --rm -v $WORK:/work "$image" python "$@"
diff --git a/image_demo.py b/image_demo.py
@@ -5,6 +5,8 @@
 import os
 
 import posenet
+import posenet.converter.tfjs2tf as tfjs2tf
+import posenet.converter.tfjsdownload as tfjsdownload
 
 
 parser = argparse.ArgumentParser()
@@ -18,9 +20,12 @@
 
 def main():
 
+    model = 'posenet'
+    neuralnet = 'mobilenet_v1_100' # mobilenet_v1_100 resnet50_v1
+    model_variant = 'stride16'
+
     with tf.compat.v1.Session() as sess:
-        model_cfg, model_outputs = posenet.load_model(args.model, sess)
-        output_stride = model_cfg['output_stride']
+        output_stride, model_outputs = posenet.load_tf_model(sess, model, neuralnet, model_variant)
 
         if args.output_dir:
             if not os.path.exists(args.output_dir):
@@ -34,10 +39,14 @@ def main():
             input_image, draw_image, output_scale = posenet.read_imgfile(
                 f, scale_factor=args.scale_factor, output_stride=output_stride)
 
+            model_cfg = tfjsdownload.model_config(model, neuralnet, model_variant)
+            input_tensor_name = model_cfg['input_tensors'][0]
+
+            # ORDER OF THE FEATURES IS DEPENDENT ON THE config.yaml file output_tensors list!!!
             heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = sess.run(
-                model_outputs,
-                feed_dict={'image:0': input_image}
-            )
+                    model_outputs,
+                    feed_dict={input_tensor_name: input_image}
+                )
 
             pose_scores, keypoint_scores, keypoint_coords = posenet.decode_multiple_poses(
                 heatmaps_result.squeeze(axis=0),
diff --git a/posenet/__init__.py b/posenet/__init__.py
@@ -1,4 +1,5 @@
 from posenet.constants import *
 from posenet.decode_multi import decode_multiple_poses
 from posenet.model import load_model
+from posenet.model import load_tf_model
 from posenet.utils import *
diff --git a/posenet/converter/common.py b/posenet/converter/common.py
@@ -0,0 +1,6 @@
+TFJS_OP_KEY = 'op'
+TFJS_DILATIONS_KEY = 'dilations'
+TFJS_CONV2D_KEY = 'Conv2D'
+
+TF_NHWC = 'NHWC'
+TF_NCHW = 'NCHW'
diff --git a/posenet/converter/config.yaml b/posenet/converter/config.yaml
@@ -1,3 +1,74 @@
+models:
+  tfjs:
+    bodypix:
+      resnet50_v1:
+        base_url: 'https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/resnet50/float'
+        model_variant:
+          stride16:
+            filename: 'model-stride16.json'
+            output_stride: 16
+            data_format: 'NHWC'
+        input_tensors:
+          - 'sub_2:0'
+        output_tensors:
+          - 'float_heatmaps:0'
+          - 'float_short_offsets:0'
+          - 'resnet_v1_50/displacement_fwd_2/BiasAdd:0'
+          - 'resnet_v1_50/displacement_bwd_2/BiasAdd:0'
+          - 'float_part_heatmaps:0'
+          - 'float_part_offsets:0'
+          - 'float_long_offsets:0'
+          - 'float_segments:0'
+      mobilenet_v1_100:
+        base_url: 'https://storage.googleapis.com/tfjs-models/savedmodel/bodypix/mobilenet/float/100'
+        model_variant:
+          stride16:
+            filename: 'model-stride16.json'
+            output_stride: 16
+            data_format: 'NHWC'
+        input_tensors:
+          - 'sub_2:0'
+        output_tensors:
+          - 'float_heatmaps:0'
+          - 'float_short_offsets:0'
+          - 'MobilenetV1/displacement_fwd_2/BiasAdd:0'
+          - 'MobilenetV1/displacement_bwd_2/BiasAdd:0'
+          - 'float_part_heatmaps:0'
+          - 'float_part_offsets:0'
+          - 'float_long_offsets:0'
+          - 'float_segments:0'
+    posenet:
+      resnet50_v1:
+        base_url: 'https://storage.googleapis.com/tfjs-models/savedmodel/posenet/resnet50/float'
+        model_variant:
+          stride16:
+            filename: 'model-stride16.json'
+            output_stride: 16
+            data_format: 'NHWC'
+        input_tensors:
+          - 'sub_2:0'
+        output_tensors:
+          - 'float_heatmaps:0'
+          - 'float_short_offsets:0'
+          - 'resnet_v1_50/displacement_fwd_2/BiasAdd:0'
+          - 'resnet_v1_50/displacement_bwd_2/BiasAdd:0'
+      mobilenet_v1_100:
+        base_url: 'https://storage.googleapis.com/tfjs-models/savedmodel/posenet/mobilenet/float/100'
+        model_variant:
+          stride16:
+            filename: 'model-stride16.json'
+            output_stride: 16
+            data_format: 'NHWC'
+        input_tensors:
+          - 'sub_2:0'
+        output_tensors:
+          - 'MobilenetV1/heatmap_2/BiasAdd:0'
+          - 'MobilenetV1/offset_2/BiasAdd:0'
+          - 'MobilenetV1/displacement_fwd_2/BiasAdd:0'
+          - 'MobilenetV1/displacement_bwd_2/BiasAdd:0'
+
+
+
 chk: 3 # 3=mobilenet_v1_101
 imageSize: 513
 GOOGLE_CLOUD_STORAGE_DIR: 'https://storage.googleapis.com/tfjs-models/weights/posenet/'
diff --git a/posenet/converter/tfjs2tf.py b/posenet/converter/tfjs2tf.py
@@ -1,8 +1,30 @@
-import tfjs_graph_converter as tfjs
 import os
+import tensorflow as tf
+import tfjs_graph_converter as tfjs
+import posenet.converter.tfjsdownload as tfjsdownload
+
+
+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, model_cfg['data_format'])
+    tfjs.api.graph_model_to_saved_model(model_cfg['tfjs_dir'], model_cfg['tf_dir'], ['serve'])
+
 
+def list_tensors(model, neuralnet, model_variant):
+    model_cfg = tfjsdownload.model_config(model, neuralnet, model_variant)
+    graph = tfjs.api.load_graph_model(model_cfg['tfjs_dir'])
+    with tf.compat.v1.Session(graph=graph) as sess:
+        # the module provides some helpers for querying model properties
+        input_tensor_names = tfjs.util.get_input_tensors(graph)
+        output_tensor_names = tfjs.util.get_output_tensors(graph)
 
-def convert(model_id, model_dir, check=False):
-    tfjsdir = os.path.join('/opt/project/_posenet_weights', 'mobilenet_v1_101')
-    tfdir = os.path.join('/opt/project/_models', 'model-mobilenet_v1_101_test')
-    tfjs.api.graph_model_to_saved_model(tfjsdir, tfdir, ['serve'])
+        print('input tensors:')
+        for it in input_tensor_names:
+            print(it)
+        print('--')
+        print('output tensors:')
+        for ot in output_tensor_names:
+            print(ot)
diff --git a/posenet/converter/tfjsdownload.py b/posenet/converter/tfjsdownload.py
@@ -0,0 +1,133 @@
+import urllib.request
+import posixpath
+import json
+import zlib
+import os
+import shutil
+import tensorflowjs.converters.common as tfjs_common
+import tfjs_graph_converter.common as tfjs_converter_common
+import posenet.converter.common as common
+
+from posenet.converter.config import load_config
+
+TFJS_MODEL_DIR = './_tfjs_models'
+TF_MODEL_DIR = './_tf_models'
+
+
+def model_config(model, neuralnet, model_variant):
+    config = load_config()
+    tfjs_models = config['models']['tfjs']
+    model_cfg = tfjs_models[model][neuralnet]
+    return {
+        'base_url': model_cfg['base_url'],
+        'filename': model_cfg['model_variant'][model_variant]['filename'],
+        'output_stride': model_cfg['model_variant'][model_variant]['output_stride'],
+        'data_format': model_cfg['model_variant'][model_variant]['data_format'],
+        'input_tensors': model_cfg['input_tensors'],
+        'output_tensors': model_cfg['output_tensors'],
+        'tfjs_dir': os.path.join(TFJS_MODEL_DIR, model, neuralnet, model_variant),
+        'tf_dir': os.path.join(TF_MODEL_DIR, model, neuralnet, model_variant)
+    }
+
+
+def _fix_dilations(node, data_format):
+    attr_key = tfjs_converter_common.TFJS_NODE_ATTR_KEY
+    dilations_key = common.TFJS_DILATIONS_KEY
+
+    if attr_key in node and dilations_key in node[attr_key]:
+        dilations = node[attr_key][dilations_key]
+        values = dilations['list']['i']
+        if len(values) == 4 and (int(values[0]) > 1 or int(values[1]) > 1) and int(values[2]) == 1 and int(values[3]) == 1:
+            if data_format == common.TF_NHWC:
+                dilation_h = values[0]
+                dilation_w = values[1]
+                values[0] = '1'
+                values[1] = dilation_h
+                values[2] = dilation_w
+                values[3] = '1'
+            if data_format == common.TF_NCHW:
+                dilation_h = values[0]
+                dilation_w = values[1]
+                values[0] = '1'
+                values[1] = '1'
+                values[2] = dilation_h
+                values[3] = dilation_w
+
+
+def _fix_model(json_model_def, data_format):
+    """
+    Input shapes are usually NHWC or NCHW and the corresponding dilations allowed are [1,x,y,1] and [1,1,x,y].
+    Some networks, like ResNet50_v1 have dilations for their Conv2D operations that are [2,2,1,1] in the model json
+    file. This representation is wrong and should be [1,2,2,1] for NHWC input and [1,1,2,2] for NCHW input.
+
+    :param data_format:
+    :param json_model_def: The JSON model definition.
+    :return: The fixed JSON model definition.
+    """
+    nodes = json_model_def[tfjs_common.ARTIFACT_MODEL_TOPOLOGY_KEY][tfjs_converter_common.TFJS_NODE_KEY]
+    for node in nodes:
+        if common.TFJS_OP_KEY in node and node[common.TFJS_OP_KEY] == common.TFJS_CONV2D_KEY:
+            _fix_dilations(node, data_format)
+
+    return json_model_def
+
+
+def fix_model_file(model_cfg):
+    model_file_path = os.path.join(model_cfg['tfjs_dir'], model_cfg['filename'])
+    with open(model_file_path, 'r') as f:
+        json_model_def = json.load(f)
+
+    json_model_def = _fix_model(json_model_def, model_cfg['data_format'])
+
+    with open(model_file_path, 'w') as f:
+        json.dump(json_model_def, f)
+
+    if not model_cfg['filename'] == 'model.json':
+        # The expected filename for the model json file is 'model.json'.
+        # See tfjs_common.ARTIFACT_MODEL_JSON_FILE_NAME in the tensorflowjs codebase.
+        normalized_model_json_file = os.path.join(model_cfg['tfjs_dir'], 'model.json')
+        shutil.copyfile(model_file_path, normalized_model_json_file)
+
+    return json_model_def
+
+
+def download_single_file(base_url, filename, save_dir):
+    output_path = os.path.join(save_dir, filename)
+    url = posixpath.join(base_url, filename)
+    req = urllib.request.Request(url)
+    response = urllib.request.urlopen(req)
+    if response.info().get('Content-Encoding') == 'gzip':
+        data = zlib.decompress(response.read(), zlib.MAX_WBITS | 32)
+    else:
+        # this path not tested since gzip encoding default on google server
+        # may need additional encoding/text handling if hit in the future
+        data = response.read()
+    with open(output_path, 'wb') as f:
+        f.write(data)
+
+
+def download_tfjs_model(model, neuralnet, model_variant, data_format):
+    """
+    Download a tfjs model with saved weights.
+
+    :param data_format:
+    :param model: The model, e.g. 'bodypix'
+    :param neuralnet: The neural net used, e.g. 'resnet50'
+    :param model_variant: The reference to the model file, e.g. 'stride16'
+    """
+    model_cfg = model_config(model, neuralnet, model_variant)
+    model_file_path = os.path.join(model_cfg['tfjs_dir'], model_cfg['filename'])
+    if os.path.exists(model_file_path):
+        print('Model file already exists: %s...' % model_file_path)
+        return
+    if not os.path.exists(model_cfg['tfjs_dir']):
+        os.makedirs(model_cfg['tfjs_dir'])
+
+    download_single_file(model_cfg['base_url'], model_cfg['filename'], model_cfg['tfjs_dir'])
+
+    json_model_def = fix_model_file(model_cfg)
+
+    shard_paths = json_model_def['weightsManifest'][0]['paths']
+    for shard in shard_paths:
+        download_single_file(model_cfg['base_url'], shard, model_cfg['tfjs_dir'])
+
diff --git a/posenet/model.py b/posenet/model.py
diff --git a/posenet/utils.py b/posenet/utils.py
diff --git a/test_tfjs2tf.py b/test_tfjs2tf.py
diff --git a/webcam_demo.py b/webcam_demo.py
