From f9b2c25fe51534df678c5f40f12d55eb49a9cf4f Mon Sep 17 00:00:00 2001
From: Tobias Eidelpes <tobias@eidelpes.info>
Date: Thu, 19 Jan 2023 21:36:58 +0100
Subject: [PATCH] Move jetson deployment code

---
 .../model.py                                  | 39 ++++++++++++++
 jetson-deployment/server.py                   | 52 +++++++++++++++++++
 2 files changed, 91 insertions(+)
 rename {yolo-second-run => jetson-deployment}/model.py (76%)
 create mode 100644 jetson-deployment/server.py

diff --git a/yolo-second-run/model.py b/jetson-deployment/model.py
similarity index 76%
rename from yolo-second-run/model.py
rename to jetson-deployment/model.py
index 54b61ac..4e2b9c5 100644
--- a/yolo-second-run/model.py
+++ b/jetson-deployment/model.py
@@ -116,11 +116,50 @@ def get_cutout(img, xmin, xmax, ymin, ymax):
     return cropped_image
 
 
+def export_to_onnx(yolo_path: str, resnet_path: str):
+    """Export the models to onnx.
+
+    :param yolo_path: path to yolo weights
+    :param resnet_path: path to resnet weights
+    :returns: None
+
+    """
+    (first, second) = load_models(yolo_path, resnet_path)
+    first.eval()
+    second.eval()
+
+    first_x = torch.randn((1, 3, 640, 640), requires_grad=True)
+    second_x = torch.randn((1, 3, 224, 224), requires_grad=True)
+
+    torch.onnx.export(first,
+                      first_x,
+                      'yolo.onnx',
+                      export_params=True,
+                      do_constant_folding=True,
+                      input_names=['input'],
+                      output_names=['output'])
+
+    torch.onnx.export(second,
+                      second_x,
+                      'resnet.onnx',
+                      export_params=True,
+                      do_constant_folding=True,
+                      input_names=['input'],
+                      output_names=['output'])
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--source', type=str, help='image file or webcam')
+    parser.add_argument('--onnx',
+                        action='store_true',
+                        dest='onnx',
+                        help='export models to onnx')
     opt = parser.parse_args()
 
     if opt.source:
         detect(opt.source, 'runs/train/yolov7-custom7/weights/best.pt',
                'resnet.pt')
+    if opt.onnx:
+        export_to_onnx('runs/train/yolov7-custom7/weights/best.pt',
+                       'resnet.pt')
diff --git a/jetson-deployment/server.py b/jetson-deployment/server.py
new file mode 100644
index 0000000..e3e3098
--- /dev/null
+++ b/jetson-deployment/server.py
@@ -0,0 +1,52 @@
+import atexit
+import jetson.utils
+
+from flask import Flask
+from apscheduler.schedulers.background import BackgroundScheduler
+
+from multiprocessing import Manager
+
+from model import detect
+
+app = Flask(__name__)
+scheduler = BackgroundScheduler(daemon=True)
+manager = Manager()
+pred = manager.dict()
+
+
+@scheduler.task('interval', id='get_pred', minutes=30, misfire_grace_time=900)
+def get_pred():
+    img = take_image('./current_image.jpg')
+    print('Job 1 executed')
+
+
+def take_image(img_path: str):
+    """Take an image with the webcam and save it to the specified
+    path.
+
+    :param str img_path: path image should be saved to
+    :returns: captured image
+
+    """
+    input = jetson.utils.videoSource('csi://0')
+    output = jetson.utils.videoOutput(img_path)
+    img = input.Capture()
+    output.Render(img)
+    return img
+
+
+
+@app.route('/')
+def index():
+    # TODO: call script and save initial image with bounding boxes
+    # TODO: get predictions and output them in JSON via API
+    # TODO: periodically get image from webcam and go to beginning
+    # TODO: JSON format: [Nr, state (1-10), timestamp, time since below 3]
+    return 'Server works'
+
+
+if __name__ == '__main__':
+    scheduler.add_job(func=get_pred, trigger='interval', minutes=30)
+    scheduler.start()
+    atexit.register(scheduler.shutdown())
+    app.run()