From 3ba91f09f7d981514b18b961ba249cc2daa831e6 Mon Sep 17 00:00:00 2001
From: YS <ys@gm.com>
Date: Tue, 21 Dec 2021 22:31:58 +0300
Subject: [PATCH] add converter to generate ssd_mobilenet_v2_coco

---
 .../ssd_mobilenet_v2_coco/Dockerfile.l4t.tf15 | 10 ++++
 converters/ssd_mobilenet_v2_coco/README.md    | 14 +++++
 .../assets/0001-fix-trt.patch                 | 52 +++++++++++++++++++
 .../ssd_mobilenet_v2_coco/assets/run.sh       |  6 +++
 converters/ssd_mobilenet_v2_coco/build.sh     | 15 ++++++
 5 files changed, 97 insertions(+)
 create mode 100644 converters/ssd_mobilenet_v2_coco/Dockerfile.l4t.tf15
 create mode 100644 converters/ssd_mobilenet_v2_coco/README.md
 create mode 100644 converters/ssd_mobilenet_v2_coco/assets/0001-fix-trt.patch
 create mode 100755 converters/ssd_mobilenet_v2_coco/assets/run.sh
 create mode 100755 converters/ssd_mobilenet_v2_coco/build.sh

diff --git a/converters/ssd_mobilenet_v2_coco/Dockerfile.l4t.tf15 b/converters/ssd_mobilenet_v2_coco/Dockerfile.l4t.tf15
new file mode 100644
index 000000000..0eac9ce6d
--- /dev/null
+++ b/converters/ssd_mobilenet_v2_coco/Dockerfile.l4t.tf15
@@ -0,0 +1,10 @@
+FROM nvcr.io/nvidia/l4t-tensorflow:r32.6.1-tf1.15-py3
+
+RUN apt-get update && apt-get install -y git sudo
+RUN git clone https://github.com/jkjung-avt/tensorrt_demos.git /tensorrt_demos
+
+ADD 0001-fix-trt.patch /tensorrt_demos/0001-fix-trt.patch
+RUN cd /tensorrt_demos && \
+    git apply 0001-fix-trt.patch
+
+ADD run.sh /run.sh
diff --git a/converters/ssd_mobilenet_v2_coco/README.md b/converters/ssd_mobilenet_v2_coco/README.md
new file mode 100644
index 000000000..188b17b4e
--- /dev/null
+++ b/converters/ssd_mobilenet_v2_coco/README.md
@@ -0,0 +1,14 @@
+
+A build.sh file will convert pre-trained tensorflow Single-Shot Multibox Detector (SSD) models through UFF to TensorRT engine to do real-time object detection with the TensorRT engine.
+
+Output will be copied to the ./model folder
+
+
+Note:
+
+This will consume pretty significant amound of memory. You might consider extending swap on Jetson Nano
+
+Usage:
+
+cd ./frigate/converters/ssd_mobilenet_v2_coco/
+./build.sh
\ No newline at end of file
diff --git a/converters/ssd_mobilenet_v2_coco/assets/0001-fix-trt.patch b/converters/ssd_mobilenet_v2_coco/assets/0001-fix-trt.patch
new file mode 100644
index 000000000..a436cb01b
--- /dev/null
+++ b/converters/ssd_mobilenet_v2_coco/assets/0001-fix-trt.patch
@@ -0,0 +1,52 @@
+From 40953eaae8ca55838e046325b257faaff0bbe33f Mon Sep 17 00:00:00 2001
+From: YS <ys@gm.com>
+Date: Tue, 21 Dec 2021 21:01:35 +0300
+Subject: [PATCH] fix trt
+
+---
+ ssd/build_engine.py | 13 ++++++++-----
+ 1 file changed, 8 insertions(+), 5 deletions(-)
+
+diff --git a/ssd/build_engine.py b/ssd/build_engine.py
+index 65729a9..e4a55c8 100644
+--- a/ssd/build_engine.py
++++ b/ssd/build_engine.py
+@@ -17,7 +17,6 @@ import uff
+ import tensorrt as trt
+ import graphsurgeon as gs
+ 
+-
+ DIR_NAME = os.path.dirname(__file__)
+ LIB_FILE = os.path.abspath(os.path.join(DIR_NAME, 'libflattenconcat.so'))
+ MODEL_SPECS = {
+@@ -286,19 +285,23 @@ def main():
+         text=True,
+         debug_mode=DEBUG_UFF)
+     with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
+-        builder.max_workspace_size = 1 << 28
++        config = builder.create_builder_config()
++        config.max_workspace_size = 1 << 28
+         builder.max_batch_size = 1
+-        builder.fp16_mode = True
++        config.set_flag(trt.BuilderFlag.FP16)
+ 
+         parser.register_input('Input', INPUT_DIMS)
+         parser.register_output('MarkOutput_0')
+         parser.parse(spec['tmp_uff'], network)
+-        engine = builder.build_cuda_engine(network)
++
++        plan = builder.build_serialized_network(network, config)
++
++        with trt.Runtime(TRT_LOGGER) as runtime:
++            engine = runtime.deserialize_cuda_engine(plan)
+ 
+         buf = engine.serialize()
+         with open(spec['output_bin'], 'wb') as f:
+             f.write(buf)
+ 
+-
+ if __name__ == '__main__':
+     main()
+-- 
+2.17.1
+
diff --git a/converters/ssd_mobilenet_v2_coco/assets/run.sh b/converters/ssd_mobilenet_v2_coco/assets/run.sh
new file mode 100755
index 000000000..28f09ad77
--- /dev/null
+++ b/converters/ssd_mobilenet_v2_coco/assets/run.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+set -xe
+cd /tensorrt_demos/ssd
+./install.sh
+python3 build_engine.py ssd_mobilenet_v2_coco
+cp /tensorrt_demos/ssd/TRT_ssd_mobilenet_v2_coco.bin /model/TRT_ssd_mobilenet_v2_coco.bin
\ No newline at end of file
diff --git a/converters/ssd_mobilenet_v2_coco/build.sh b/converters/ssd_mobilenet_v2_coco/build.sh
new file mode 100755
index 000000000..8f259c5a1
--- /dev/null
+++ b/converters/ssd_mobilenet_v2_coco/build.sh
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+mkdir -p $(pwd)/model
+
+docker build --tag models.ssd_v2_coco  --file ./Dockerfile.l4t.tf15 ./assets/
+
+sudo docker run --rm -it --name models.ssd_v2_coco \
+    --mount type=tmpfs,target=/tmp/cache,tmpfs-size=1000000000 \
+    -v $(pwd)/model:/model:rw \
+    -v /tmp/argus_socket:/tmp/argus_socket \
+    -e NVIDIA_VISIBLE_DEVICES=all \
+    -e NVIDIA_DRIVER_CAPABILITIES=compute,utility,video \
+    --runtime=nvidia \
+    --privileged \
+    models.ssd_v2_coco /run.sh