Update jetson-nano model to Jetpack 4.5

README.md

@@ -127,7 +127,7 @@ git checkout $(git tag | tail -1)
 ```
 
 ###### Run on Jetson Nano
-* You need to have JetPack 4.3 installed on your Jetson Nano.
+* You need to have JetPack 4.5 installed on your Jetson Nano.
 
 ```bash
 # 1) Download TensorRT engine file built with JetPack 4.3:
@@ -243,7 +243,7 @@ Before running any of the images available in the Docker repository, you need to
 Alternatively you may simply pull the folder structure from this repository.
 
 ###### Run on Jetson Nano
-* You need to have JetPack 4.3 installed on your Jetson Nano.
+* You need to have JetPack 4.5 installed on your Jetson Nano.
 ```bash
 # Download TensorRT engine file built with JetPack 4.3:
 mkdir data/jetson
@@ -326,7 +326,7 @@ docker run -it -p HOST_PORT:8000 -v $PWD/data:/repo/data -v $PWD/config-x86-open
 ##### Running the processor with docker-compose
 
 ###### Run on Jetson Nano
-* You need to have JetPack 4.3 installed on your Jetson Nano.
+* You need to have JetPack 4.5 installed on your Jetson Nano.
 
 ```bash
 # 1) Download TensorRT engine file built with JetPack 4.3:

api/requirements.txt

@@ -1,7 +1,7 @@
 aiofiles==0.5.0
 boto3==1.14.59
 fastapi==0.61.1
-pandas==1.1.2
+pandas==1.1.5
 pyhumps==1.6.1
 pytest==6.0.1
 pytest-mock==3.5.1

exporters/jetson/add_plugin_and_preprocess_ssd_mobilenet.py

@@ -0,0 +1,164 @@
+"""
+Add plugin nodes for custom layers of ssd-mobilenet-v1/v2 to the graph
+and replace some newer operations which existing previous ones.
+
+Reference:
+   https://github.com/jkjung-avt/tensorrt_demos/blob/master/ssd/build_engine.py
+"""
+
+import tensorrt as trt
+import graphsurgeon as gs
+import tensorflow as tf
+import numpy as np
+
+def replace_addv2(graph):
+    """Replace all 'AddV2' in the graph with 'Add'.
+    NOTE: 'AddV2' is not supported by UFF parser.
+    """
+    for node in graph.find_nodes_by_op('AddV2'):
+        gs.update_node(node, op='Add')
+    return graph
+
+
+def replace_fusedbnv3(graph):
+    """Replace all 'FusedBatchNormV3' in the graph with 'FusedBatchNorm'.
+    NOTE: 'FusedBatchNormV3' is not supported by UFF parser.
+    https://devtalk.nvidia.com/default/topic/1066445/tensorrt/tensorrt-6-0-1-tensorflow-1-14-no-conversion-function-registered-for-layer-fusedbatchnormv3-yet/post/5403567/#5403567
+    """
+    for node in graph.find_nodes_by_op('FusedBatchNormV3'):
+        gs.update_node(node, op='FusedBatchNorm')
+    return graph
+
+def parse_gridAnchor(graph):
+    """
+    define a constant input tensor and set that as the input for the GridAnchor node 
+    as UFF file does not provide an input element for the GridAnchor node
+    """
+
+    data = np.array([1, 1], dtype=np.float32) 
+    anchor_input = gs.create_node("AnchorInput", "Const", value=data, dtype=tf.float32)  
+    graph.append(anchor_input)
+    graph.find_nodes_by_op("GridAnchor_TRT")[0].input.insert(0, "AnchorInput")
+
+    return graph
+
+
+
+def add_plugin_and_preprocess(graph, model, num_classes, neuralet_adaptive_model):
+    """add_plugin
+
+    Reference:
+    1. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v1_coco_2018_01_28.py
+    2. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v2_coco_2018_03_29.py
+    3. https://devtalk.nvidia.com/default/topic/1050465/jetson-nano/how-to-write-config-py-for-converting-ssd-mobilenetv2-to-uff-format/post/5333033/#5333033
+    """
+    num_classes += 1
+    min_size = 0.2
+    max_size = 0.95
+    input_order = [0, 2, 1]
+    if neuralet_adaptive_model == 0:
+        input_order = [1, 0, 2]
+    input_dims = (3, 300, 300)
+
+    all_assert_nodes = graph.find_nodes_by_op("Assert")
+    graph.remove(all_assert_nodes, remove_exclusive_dependencies=True)
+
+    all_identity_nodes = graph.find_nodes_by_op("Identity")
+    graph.forward_inputs(all_identity_nodes)
+    Input = gs.create_plugin_node(
+        name="Input",
+        op="Placeholder",
+        shape=(1,) + input_dims,
+        dtype=tf.float32
+    )
+
+    PriorBox = gs.create_plugin_node(
+        name="MultipleGridAnchorGenerator",
+        op="GridAnchor_TRT",
+        minSize=min_size,  # was 0.2
+        maxSize=max_size,  # was 0.95
+        aspectRatios=[1.0, 2.0, 0.5, 3.0, 0.33],
+        variance=[0.1, 0.1, 0.2, 0.2],
+        featureMapShapes=[19, 10, 5, 3, 2, 1],
+        numLayers=6,
+        dtype=tf.float32
+    )
+
+    NMS = gs.create_plugin_node(
+        name="NMS",
+        op="NMS_TRT",
+        shareLocation=1,
+        varianceEncodedInTarget=0,
+        backgroundLabelId=0,
+        confidenceThreshold=0.3,  # was 1e-8
+        nmsThreshold=0.6,
+        topK=100,
+        keepTopK=100,
+        numClasses=num_classes,  # was 91
+        inputOrder=input_order,
+        confSigmoid=1,
+        isNormalized=1,
+        dtype=tf.float32
+    )
+
+    concat_priorbox = gs.create_node(
+        "concat_priorbox",
+        op="ConcatV2",
+        axis=2,
+        dtype=tf.float32
+    )
+    if trt.__version__[0] >= '6':
+        concat_box_loc = gs.create_plugin_node(
+            "concat_box_loc",
+            op="FlattenConcat_TRT",
+            axis=1,
+            ignoreBatch=0,
+            dtype=tf.float32
+        )
+        concat_box_conf = gs.create_plugin_node(
+            "concat_box_conf",
+            op="FlattenConcat_TRT",
+            axis=1,
+            ignoreBatch=0,
+            dtype=tf.float32
+        )
+    else:
+        concat_box_loc = gs.create_plugin_node(
+            "concat_box_loc",
+            op="FlattenConcat_TRT",
+            dtype=tf.float32
+        )
+        concat_box_conf = gs.create_plugin_node(
+            "concat_box_conf",
+            op="FlattenConcat_TRT",
+            dtype=tf.float32
+        )
+
+    namespace_plugin_map = {
+        "MultipleGridAnchorGenerator": PriorBox,
+        "Postprocessor": NMS,
+        "Preprocessor": Input,
+        "Cast": Input,
+        "ToFloat": Input,
+        "image_tensor": Input,
+        "MultipleGridAnchorGenerator/Concatenate": concat_priorbox,  # for 'ssd_mobilenet_v1_coco'
+        "Concatenate": concat_priorbox, 
+        "concat": concat_box_loc,
+        "concat_1": concat_box_conf
+    }
+
+    graph.collapse_namespaces(namespace_plugin_map)
+
+    graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False)
+    graph.find_nodes_by_op("NMS_TRT")[0].input.remove("Input")
+
+    if ( "image_tensor:0" in graph.find_nodes_by_name("Input")[0].input ):
+    # for ssd_mobilenet_v1
+        graph.find_nodes_by_name("Input")[0].input.remove("image_tensor:0")
+
+    graph = replace_addv2(graph)
+    graph = replace_fusedbnv3(graph)
+    graph = parse_gridAnchor(graph)
+
+
+    return graph

exporters/jetson/generate_mobilenet_tensorrt.bash

@@ -0,0 +1,27 @@
+#!/bin/bash
+# This script generates a TensorRT engine from pretrained Tensorflow SSD Mobilenet v2 COCO model.
+
+echo $relative_path
+if [ $# -eq 2 ]
+then
+	pb_file=$1
+	adaptive_model=1
+	num_classes=$2
+
+else 
+	if [ ! -f "$relative_path/data/jetson/ssd_mobilenet_v2_coco_2018_03_29.tar.gz"  ]; then
+		echo "downloading ssd_mobilenet_v2_coco_2018_03_29 ..."
+		wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_coco_2018_03_29.tar.gz -P /repo$relative_path/data/jetson/
+	fi
+	tar -xvf /repo$relative_path/data/jetson/ssd_mobilenet_v2_coco_2018_03_29.tar.gz --no-same-owner -C /repo$relative_path/data/jetson/
+	cp "/repo$relative_path/data/jetson/ssd_mobilenet_v2_coco_2018_03_29/frozen_inference_graph.pb" "/repo$relative_path/data/jetson/ssd_mobilenet_v2_coco_2018_03_29/TRT_ssd_mobilenet_v2_coco.pb"
+	rm "/repo$relative_path/data/jetson/ssd_mobilenet_v2_coco_2018_03_29/frozen_inference_graph.pb"
+	pb_file="/repo$relative_path/data/jetson/ssd_mobilenet_v2_coco_2018_03_29/TRT_ssd_mobilenet_v2_coco.pb"
+	adaptive_model=0
+	num_classes=90
+
+fi
+
+
+echo "************  Generating TensorRT from: $pb_file  **************"
+python3 /repo$relative_path/exporters/jetson/trt_exporter.py --pb_file $pb_file --out_dir /repo$relative_path/data/jetson --neuralet_adaptive_model $adaptive_model --num_classes $num_classes

exporters/jetson/trt_exporter.py

@@ -0,0 +1,97 @@
+"""build_engine.py
+
+This script converts a SSD model (pb) to UFF and subsequently builds
+the TensorRT engine.
+
+Input : spces of a ssd frozen inference graph in config.ini file
+Output: TensorRT Engine file
+
+Reference:
+   https://github.com/jkjung-avt/tensorrt_demos/blob/master/ssd/build_engine.py
+"""
+
+
+import os
+import ctypes
+import argparse
+import wget
+
+import uff
+import tensorrt as trt
+import graphsurgeon as gs
+import numpy as np
+import add_plugin_and_preprocess_ssd_mobilenet as plugin
+
+def export_trt(pb_file, output_dir, num_classes=1, neuralet_adaptive_model=1):
+    """
+    Exports the Tensorflow pb models to TensorRT engines.
+    Args:
+        pb_file: The path of input pb file
+        output_dir: A directory to store the output files
+        num_classes: Detector's number of classes
+    """
+    lib_flatten_concat_file = "exporters/jetson/libflattenconcat.so.6"
+    # initialize
+    if trt.__version__[0] < '7':
+        ctypes.CDLL(lib_flatten_concat_file)
+    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+    trt.init_libnvinfer_plugins(TRT_LOGGER, '')
+
+    # compile the model into TensorRT engine
+    model = "ssd_mobilenet_v2_coco"
+
+    if not os.path.isfile(pb_file):
+        raise FileNotFoundError('model does not exist under: {}'.format(pb_file))
+
+    if not os.path.isdir(output_dir):
+        print("the provided output directory : {0} is not exist".format(output_dir))
+        print("creating output directory : {0}".format(output_dir))
+        os.makedirs(output_dir, exist_ok=True)
+
+
+    dynamic_graph = plugin.add_plugin_and_preprocess(
+        gs.DynamicGraph(pb_file),
+        model,
+        num_classes,
+        neuralet_adaptive_model)
+    model_file_name = ".".join((pb_file.split("/")[-1]).split(".")[:-1])
+    uff_path = os.path.join(output_dir, model_file_name + ".uff")
+    _ = uff.from_tensorflow(
+        dynamic_graph.as_graph_def(),
+        output_nodes=['NMS'],
+        output_filename=uff_path,
+        text=True,
+        debug_mode=False)
+    input_dims = (3, 300, 300)
+    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
+        config = builder.create_builder_config()
+        config.max_workspace_size = 1 << 28
+        builder.max_batch_size = 1
+        config.set_flag(trt.BuilderFlag.FP16)
+
+        parser.register_input('Input', input_dims)
+        parser.register_output('MarkOutput_0')
+        parser.parse(uff_path, network)
+        engine = builder.build_engine(network, config)
+
+        buf = engine.serialize()
+        engine_path = os.path.join(output_dir, model_file_name + ".bin")
+        with open(engine_path, 'wb') as f:
+            f.write(buf)
+        print("your model has been converted to trt engine successfully under : {}".format(engine_path))
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="This script exports a pb tensorflow model to a TRT engine")
+    parser.add_argument("--pb_file", type=str, required=True, help="the path of input pb file")
+    parser.add_argument("--out_dir", type=str, required=True, help="a directory to store the output files")
+    parser.add_argument("--num_classes", type=int, default=90, help="detector's number of classes")
+    parser.add_argument("--neuralet_adaptive_model", type=int, default=1, help="1 if the model is trained by Neuralet adaptive learning, 0 if not")
+    args = parser.parse_args()
+    pb_file = args.pb_file
+    output_dir = args.out_dir
+    num_classes = args.num_classes
+    neuralet_adaptive_model = args.neuralet_adaptive_model
+    export_trt(pb_file=pb_file, output_dir=output_dir, num_classes=num_classes, neuralet_adaptive_model=neuralet_adaptive_model)
+