From e724fe3da6795f71d549ee5df347889e2c330b8d Mon Sep 17 00:00:00 2001
From: Blake Blackshear <blakeb@blakeshome.com>
Date: Sun, 3 Apr 2022 15:00:11 -0500
Subject: [PATCH] add endpoint to submit to plus

---
 docker/Dockerfile                   |  3 +-
 frigate/app.py                      |  9 ++-
 frigate/const.py                    |  2 +
 frigate/http.py                     | 62 ++++++++++++++++---
 frigate/models.py                   |  1 +
 frigate/plus.py                     | 96 +++++++++++++++++++++++++++++
 migrations/010_add_plus_image_id.py | 46 ++++++++++++++
 7 files changed, 210 insertions(+), 9 deletions(-)
 create mode 100644 frigate/plus.py
 create mode 100644 migrations/010_add_plus_image_id.py

diff --git a/docker/Dockerfile b/docker/Dockerfile
index d758a8b5a..622da1432 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -57,7 +57,8 @@ RUN pip3 wheel --wheel-dir=/wheels \
     peewee_migrate \
     pydantic \
     zeroconf \
-    ws4py
+    ws4py \
+    requests
 
 # Frigate Container
 FROM debian:11-slim
diff --git a/frigate/app.py b/frigate/app.py
index bf593d6ac..8f0aa36da 100644
--- a/frigate/app.py
+++ b/frigate/app.py
@@ -16,7 +16,7 @@ from playhouse.sqliteq import SqliteQueueDatabase
 from pydantic import ValidationError
 
 from frigate.config import DetectorTypeEnum, FrigateConfig
-from frigate.const import CACHE_DIR, CLIPS_DIR, RECORD_DIR
+from frigate.const import CACHE_DIR, CLIPS_DIR, RECORD_DIR, PLUS_ENV_VAR, PLUS_API_HOST
 from frigate.edgetpu import EdgeTPUProcess
 from frigate.events import EventCleanup, EventProcessor
 from frigate.http import create_app
@@ -25,6 +25,7 @@ from frigate.models import Event, Recordings
 from frigate.mqtt import MqttSocketRelay, create_mqtt_client
 from frigate.object_processing import TrackedObjectProcessor
 from frigate.output import output_frames
+from frigate.plus import PlusApi
 from frigate.record import RecordingCleanup, RecordingMaintainer
 from frigate.stats import StatsEmitter, stats_init
 from frigate.version import VERSION
@@ -44,6 +45,11 @@ class FrigateApp:
         self.detection_out_events: Dict[str, mp.Event] = {}
         self.detection_shms: List[mp.shared_memory.SharedMemory] = []
         self.log_queue = mp.Queue()
+        self.plus_api = (
+            PlusApi(PLUS_API_HOST, os.environ.get(PLUS_ENV_VAR))
+            if PLUS_ENV_VAR in os.environ
+            else None
+        )
         self.camera_metrics = {}
 
     def set_environment_vars(self):
@@ -146,6 +152,7 @@ class FrigateApp:
             self.db,
             self.stats_tracking,
             self.detected_frames_processor,
+            self.plus_api,
         )
 
     def init_mqtt(self):
diff --git a/frigate/const.py b/frigate/const.py
index afb6075a7..004e3c28e 100644
--- a/frigate/const.py
+++ b/frigate/const.py
@@ -3,3 +3,5 @@ CLIPS_DIR = f"{BASE_DIR}/clips"
 RECORD_DIR = f"{BASE_DIR}/recordings"
 CACHE_DIR = "/tmp/cache"
 YAML_EXT = (".yaml", ".yml")
+PLUS_ENV_VAR = "PLUS_API_KEY"
+PLUS_API_HOST = "https://api.frigate.video"
diff --git a/frigate/http.py b/frigate/http.py
index f21d1f8e6..ea7054e3c 100644
--- a/frigate/http.py
+++ b/frigate/http.py
@@ -29,7 +29,7 @@ from flask import (
 from peewee import SqliteDatabase, operator, fn, DoesNotExist, Value
 from playhouse.shortcuts import model_to_dict
 
-from frigate.const import CLIPS_DIR, RECORD_DIR
+from frigate.const import CLIPS_DIR, PLUS_ENV_VAR
 from frigate.models import Event, Recordings
 from frigate.stats import stats_snapshot
 from frigate.util import calculate_region
@@ -45,6 +45,7 @@ def create_app(
     database: SqliteDatabase,
     stats_tracking,
     detected_frames_processor,
+    plus_api,
 ):
     app = Flask(__name__)
 
@@ -61,6 +62,7 @@ def create_app(
     app.frigate_config = frigate_config
     app.stats_tracking = stats_tracking
     app.detected_frames_processor = detected_frames_processor
+    app.plus_api = plus_api
 
     app.register_blueprint(bp)
 
@@ -137,6 +139,52 @@ def set_retain(id):
     )
 
 
+@bp.route("/events/<id>/plus", methods=("POST",))
+def send_to_plus(id):
+    if current_app.plus_api is None:
+        return make_response(
+            jsonify({"success": False, "message": "Plus token not set"}), 400
+        )
+
+    try:
+        event = Event.get(Event.id == id)
+    except DoesNotExist:
+        return make_response(
+            jsonify({"success": False, "message": "Event" + id + " not found"}), 404
+        )
+
+    if event.plus_id:
+        return make_response(
+            jsonify({"success": False, "message": "Already submitted to plus"}), 400
+        )
+
+    # load clean.png
+    try:
+        filename = f"{event.camera}-{event.id}-clean.png"
+        image = cv2.imread(os.path.join(CLIPS_DIR, filename))
+    except Exception:
+        return make_response(
+            jsonify(
+                {"success": False, "message": "Unable to load clean png for event"}
+            ),
+            400,
+        )
+
+    try:
+        plus_id = current_app.plus_api.upload_image(image, event.camera)
+    except Exception as ex:
+        return make_response(
+            jsonify({"success": False, "message": str(ex)}),
+            400,
+        )
+
+    # store image id in the database
+    event.plus_id = plus_id
+    event.save()
+
+    return "success"
+
+
 @bp.route("/events/<id>/retain", methods=("DELETE",))
 def delete_retain(id):
     try:
@@ -252,6 +300,7 @@ def event_thumbnail(id):
         response.headers["Cache-Control"] = "private, max-age=31536000"
     return response
 
+
 @bp.route("/<camera_name>/<label>/best.jpg")
 @bp.route("/<camera_name>/<label>/thumbnail.jpg")
 def label_thumbnail(camera_name, label):
@@ -277,9 +326,7 @@ def label_thumbnail(camera_name, label):
         return event_thumbnail(event.id)
     except DoesNotExist:
         frame = np.zeros((175, 175, 3), np.uint8)
-        ret, jpg = cv2.imencode(
-            ".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70]
-        )
+        ret, jpg = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
 
         response = make_response(jpg.tobytes())
         response.headers["Content-Type"] = "image/jpeg"
@@ -330,6 +377,7 @@ def event_snapshot(id):
         ] = f"attachment; filename=snapshot-{id}.jpg"
     return response
 
+
 @bp.route("/<camera_name>/<label>/snapshot.jpg")
 def label_snapshot(camera_name, label):
     if label == "any":
@@ -353,9 +401,7 @@ def label_snapshot(camera_name, label):
         return event_snapshot(event.id)
     except DoesNotExist:
         frame = np.zeros((720, 1280, 3), np.uint8)
-        ret, jpg = cv2.imencode(
-            ".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70]
-        )
+        ret, jpg = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
 
         response = make_response(jpg.tobytes())
         response.headers["Content-Type"] = "image/jpeg"
@@ -460,6 +506,8 @@ def config():
         for cmd in camera_dict["ffmpeg_cmds"]:
             cmd["cmd"] = " ".join(cmd["cmd"])
 
+    config["plus"] = {"enabled": PLUS_ENV_VAR in os.environ}
+
     return jsonify(config)
 
 
diff --git a/frigate/models.py b/frigate/models.py
index 70451983f..05f4f6d82 100644
--- a/frigate/models.py
+++ b/frigate/models.py
@@ -21,6 +21,7 @@ class Event(Model):
     area = IntegerField()
     retain_indefinitely = BooleanField(default=False)
     ratio = FloatField(default=1.0)
+    plus_id = CharField(max_length=30)
 
 
 class Recordings(Model):
diff --git a/frigate/plus.py b/frigate/plus.py
new file mode 100644
index 000000000..5e74e7592
--- /dev/null
+++ b/frigate/plus.py
@@ -0,0 +1,96 @@
+import datetime
+import logging
+import requests
+import cv2
+
+logger = logging.getLogger(__name__)
+
+
+def get_jpg_bytes(image, max_dim, quality):
+    if image.shape[1] >= image.shape[0]:
+        width = min(max_dim, image.shape[1])
+        height = int(width * image.shape[0] / image.shape[1])
+    else:
+        height = min(max_dim, image.shape[0])
+        width = int(height * image.shape[1] / image.shape[0])
+
+    original = cv2.resize(image, dsize=(width, height), interpolation=cv2.INTER_AREA)
+
+    ret, jpg = cv2.imencode(".jpg", original, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
+    return jpg.tobytes()
+
+
+class PlusApi:
+    def __init__(self, host, key: str):
+        self.host = host
+        self.key = key
+        self._token_data = None
+
+    def _refresh_token_if_needed(self):
+        if (
+            self._token_data is None
+            or self._token_data["expires"] - datetime.datetime.now().timestamp() < 60
+        ):
+            parts = self.key.split(":")
+            r = requests.get(f"{self.host}/v1/auth/token", auth=(parts[0], parts[1]))
+            self._token_data = r.json()
+
+    def _get_authorization_header(self):
+        self._refresh_token_if_needed()
+        return {"authorization": f"Bearer {self._token_data['accessToken']}"}
+
+    def _get(self, path):
+        return requests.get(
+            f"{self.host}/v1/{path}", headers=self._get_authorization_header()
+        )
+
+    def _post(self, path, data):
+        return requests.post(
+            f"{self.host}/v1/{path}",
+            headers=self._get_authorization_header(),
+            json=data,
+        )
+
+    def upload_image(self, image, camera: str):
+        r = self._get("image/signed_urls")
+        presigned_urls = r.json()
+        if not r.ok:
+            logger.exception(ex)
+            raise Exception("Unable to get signed urls")
+
+        # resize and submit original
+        files = {"file": get_jpg_bytes(image, 1920, 85)}
+        data = presigned_urls["original"]["fields"]
+        data["content-type"] = "image/jpeg"
+        r = requests.post(presigned_urls["original"]["url"], files=files, data=data)
+        if not r.ok:
+            logger.error(f"Failed to upload original: {r.status_code} {r.text}")
+            raise Exception(r.text)
+
+        # resize and submit annotate
+        files = {"file": get_jpg_bytes(image, 640, 70)}
+        data = presigned_urls["annotate"]["fields"]
+        data["content-type"] = "image/jpeg"
+        r = requests.post(presigned_urls["annotate"]["url"], files=files, data=data)
+        if not r.ok:
+            logger.error(f"Failed to upload annotate: {r.status_code} {r.text}")
+            raise Exception(r.text)
+
+        # resize and submit thumbnail
+        files = {"file": get_jpg_bytes(image, 200, 70)}
+        data = presigned_urls["thumbnail"]["fields"]
+        data["content-type"] = "image/jpeg"
+        r = requests.post(presigned_urls["thumbnail"]["url"], files=files, data=data)
+        if not r.ok:
+            logger.error(f"Failed to upload thumbnail: {r.status_code} {r.text}")
+            raise Exception(r.text)
+
+        # create image
+        r = self._post(
+            "image/create", {"id": presigned_urls["imageId"], "camera": camera}
+        )
+        if not r.ok:
+            raise Exception(r.text)
+
+        # return image id
+        return presigned_urls["imageId"]
diff --git a/migrations/010_add_plus_image_id.py b/migrations/010_add_plus_image_id.py
new file mode 100644
index 000000000..0e239e58b
--- /dev/null
+++ b/migrations/010_add_plus_image_id.py
@@ -0,0 +1,46 @@
+"""Peewee migrations -- 010_add_plus_image_id.py.
+
+Some examples (model - class or model name)::
+
+    > Model = migrator.orm['model_name']            # Return model in current state by name
+
+    > migrator.sql(sql)                             # Run custom SQL
+    > migrator.python(func, *args, **kwargs)        # Run python code
+    > migrator.create_model(Model)                  # Create a model (could be used as decorator)
+    > migrator.remove_model(model, cascade=True)    # Remove a model
+    > migrator.add_fields(model, **fields)          # Add fields to a model
+    > migrator.change_fields(model, **fields)       # Change fields
+    > migrator.remove_fields(model, *field_names, cascade=True)
+    > migrator.rename_field(model, old_field_name, new_field_name)
+    > migrator.rename_table(model, new_table_name)
+    > migrator.add_index(model, *col_names, unique=False)
+    > migrator.drop_index(model, *col_names)
+    > migrator.add_not_null(model, *field_names)
+    > migrator.drop_not_null(model, *field_names)
+    > migrator.add_default(model, field_name, default)
+
+"""
+
+import datetime as dt
+import peewee as pw
+from playhouse.sqlite_ext import *
+from decimal import ROUND_HALF_EVEN
+from frigate.models import Event
+
+try:
+    import playhouse.postgres_ext as pw_pext
+except ImportError:
+    pass
+
+SQL = pw.SQL
+
+
+def migrate(migrator, database, fake=False, **kwargs):
+    migrator.add_fields(
+        Event,
+        plus_id=pw.CharField(max_length=30, null=True),
+    )
+
+
+def rollback(migrator, database, fake=False, **kwargs):
+    migrator.remove_fields(Event, ["plus_id"])