From b1cc64d4fa4bcf259717dd85110b4e09d9cd92c6 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Thu, 17 Mar 2022 06:18:43 -0600
Subject: [PATCH] FEAT: Ability to set sub labels for specific events (#2949)

* Add sub label to model and set / delete funs

* Add migrations for sub label

* Tweaks to API and model

* Show sublabel if available

* Cleanups

* Update docs

* Show person in UI title

* Fix typo and don't fail on no json

* Transfer sub labels for in progress events

* Remove sublabel reset

* Remove person only check

* Make default null

* Update docs and formatting

* Make default null

* Make nullable in migration

* Undo null

* Update model to accept null

* Update migration to accept null

* Don't set to default values

* Remove redundant defaults and update http logic

* Only need a single route

* Enforce 20 character limit in http

* Update docs to mention 20 character limit

* Cleanup

* Separate insert and update to make sure updated values are retained when event ends

* Use insert instead of replace

* Remove redundant if and have should_update_db include clip or snapshot requirement.
---
 docs/docs/integrations/api.md   | 13 +++++-
 frigate/events.py               | 77 ++++++++++++++++++++++-----------
 frigate/http.py                 | 38 +++++++++++++---
 frigate/models.py               |  1 +
 migrations/008_add_sub_label.py | 46 ++++++++++++++++++++
 web/src/routes/Events.jsx       |  2 +-
 6 files changed, 144 insertions(+), 33 deletions(-)
 create mode 100644 migrations/008_add_sub_label.py

diff --git a/docs/docs/integrations/api.md b/docs/docs/integrations/api.md
index 52e6772fd..1569a6774 100644
--- a/docs/docs/integrations/api.md
+++ b/docs/docs/integrations/api.md
@@ -186,13 +186,24 @@ Sets retain to true for the event id.
 
 Sets retain to false for the event id (event may be deleted quickly after removing).
 
+### `POST /api/events/<id>/sub_label`
+
+Set a sub label for an event. For example to update `person` -> `person's name` if they were recognized with facial recognition.
+Sub labels must be 20 characters or shorter.
+
+```json
+{
+  "subLabel": "some_string"
+}
+```
+
 ### `GET /api/events/<id>/thumbnail.jpg`
 
 Returns a thumbnail for the event id optimized for notifications. Works while the event is in progress and after completion. Passing `?format=android` will convert the thumbnail to 2:1 aspect ratio.
 
 ### `GET /api/<camera_name>/<label>/thumbnail.jpg`
 
-Returns the thumbnail from the latest event for the given camera and label combo. Using `any` as the label will return the latest thumbnail regardless of type. 
+Returns the thumbnail from the latest event for the given camera and label combo. Using `any` as the label will return the latest thumbnail regardless of type.
 
 ### `GET /api/events/<id>/clip.mp4`
 
diff --git a/frigate/events.py b/frigate/events.py
index b90204361..efd4b6e63 100644
--- a/frigate/events.py
+++ b/frigate/events.py
@@ -14,14 +14,22 @@ from frigate.models import Event
 
 logger = logging.getLogger(__name__)
 
+def should_insert_db(prev_event, current_event):
+    """If current event has new clip or snapshot."""
+    return (
+        (not prev_event["has_clip"] and not prev_event["has_snapshot"])
+        and (current_event["has_clip"] or current_event["has_snapshot"])
+    )
 
 def should_update_db(prev_event, current_event):
+    """If current_event has updated fields and (clip or snapshot)."""
     return (
-        prev_event["top_score"] != current_event["top_score"]
-        or prev_event["entered_zones"] != current_event["entered_zones"]
-        or prev_event["thumbnail"] != current_event["thumbnail"]
-        or prev_event["has_clip"] != current_event["has_clip"]
-        or prev_event["has_snapshot"] != current_event["has_snapshot"]
+        (current_event["has_clip"] or current_event["has_snapshot"])
+        and (prev_event["top_score"] != current_event["top_score"]
+            or prev_event["entered_zones"] != current_event["entered_zones"]
+            or prev_event["thumbnail"] != current_event["thumbnail"]
+            or prev_event["has_clip"] != current_event["has_clip"]
+            or prev_event["has_snapshot"] != current_event["has_snapshot"])
     )
 
 
@@ -58,33 +66,52 @@ class EventProcessor(threading.Thread):
             if event_type == "start":
                 self.events_in_process[event_data["id"]] = event_data
 
+            elif event_type == "update" and should_insert_db(
+                self.events_in_process[event_data["id"]], event_data
+            ):
+                self.events_in_process[event_data["id"]] = event_data
+                # TODO: this will generate a lot of db activity possibly
+                Event.insert(
+                    id=event_data["id"],
+                    label=event_data["label"],
+                    camera=camera,
+                    start_time=event_data["start_time"] - event_config.pre_capture,
+                    end_time=None,
+                    top_score=event_data["top_score"],
+                    false_positive=event_data["false_positive"],
+                    zones=list(event_data["entered_zones"]),
+                    thumbnail=event_data["thumbnail"],
+                    region=event_data["region"],
+                    box=event_data["box"],
+                    area=event_data["area"],
+                    has_clip=event_data["has_clip"],
+                    has_snapshot=event_data["has_snapshot"],
+                ).execute()
+
             elif event_type == "update" and should_update_db(
                 self.events_in_process[event_data["id"]], event_data
             ):
                 self.events_in_process[event_data["id"]] = event_data
                 # TODO: this will generate a lot of db activity possibly
-                if event_data["has_clip"] or event_data["has_snapshot"]:
-                    Event.replace(
-                        id=event_data["id"],
-                        label=event_data["label"],
-                        camera=camera,
-                        start_time=event_data["start_time"] - event_config.pre_capture,
-                        end_time=None,
-                        top_score=event_data["top_score"],
-                        false_positive=event_data["false_positive"],
-                        zones=list(event_data["entered_zones"]),
-                        thumbnail=event_data["thumbnail"],
-                        region=event_data["region"],
-                        box=event_data["box"],
-                        area=event_data["area"],
-                        has_clip=event_data["has_clip"],
-                        has_snapshot=event_data["has_snapshot"],
-                    ).execute()
+                Event.update(
+                    label=event_data["label"],
+                    camera=camera,
+                    start_time=event_data["start_time"] - event_config.pre_capture,
+                    end_time=None,
+                    top_score=event_data["top_score"],
+                    false_positive=event_data["false_positive"],
+                    zones=list(event_data["entered_zones"]),
+                    thumbnail=event_data["thumbnail"],
+                    region=event_data["region"],
+                    box=event_data["box"],
+                    area=event_data["area"],
+                    has_clip=event_data["has_clip"],
+                    has_snapshot=event_data["has_snapshot"],
+                ).where(Event.id == event_data["id"]).execute()
 
             elif event_type == "end":
                 if event_data["has_clip"] or event_data["has_snapshot"]:
-                    Event.replace(
-                        id=event_data["id"],
+                    Event.update(
                         label=event_data["label"],
                         camera=camera,
                         start_time=event_data["start_time"] - event_config.pre_capture,
@@ -98,7 +125,7 @@ class EventProcessor(threading.Thread):
                         area=event_data["area"],
                         has_clip=event_data["has_clip"],
                         has_snapshot=event_data["has_snapshot"],
-                    ).execute()
+                    ).where(Event.id == event_data["id"]).execute()
 
                 del self.events_in_process[event_data["id"]]
                 self.event_processed_queue.put((event_data["id"], camera))
diff --git a/frigate/http.py b/frigate/http.py
index dfa4b5f62..f21d1f8e6 100644
--- a/frigate/http.py
+++ b/frigate/http.py
@@ -126,14 +126,14 @@ def set_retain(id):
         event = Event.get(Event.id == id)
     except DoesNotExist:
         return make_response(
-            jsonify({"success": False, "message": "Event" + id + " not found"}), 404
+            jsonify({"success": False, "message": "Event " + id + " not found"}), 404
         )
 
     event.retain_indefinitely = True
     event.save()
 
     return make_response(
-        jsonify({"success": True, "message": "Event" + id + " retained"}), 200
+        jsonify({"success": True, "message": "Event " + id + " retained"}), 200
     )
 
 
@@ -143,16 +143,42 @@ def delete_retain(id):
         event = Event.get(Event.id == id)
     except DoesNotExist:
         return make_response(
-            jsonify({"success": False, "message": "Event" + id + " not found"}), 404
+            jsonify({"success": False, "message": "Event " + id + " not found"}), 404
         )
 
     event.retain_indefinitely = False
     event.save()
 
     return make_response(
-        jsonify({"success": True, "message": "Event" + id + " un-retained"}), 200
+        jsonify({"success": True, "message": "Event " + id + " un-retained"}), 200
     )
 
+@bp.route("/events/<id>/sub_label", methods=("POST",))
+def set_sub_label(id):
+    try:
+        event = Event.get(Event.id == id)
+    except DoesNotExist:
+        return make_response(
+            jsonify({"success": False, "message": "Event " + id + " not found"}), 404
+        )
+
+    if request.json:
+        new_sub_label = request.json.get("subLabel")
+    else:
+        new_sub_label = None
+
+
+    if new_sub_label and len(new_sub_label) > 20:
+        return make_response(
+            jsonify({"success": False, "message": new_sub_label + " exceeds the 20 character limit for sub_label"}), 400
+        )
+
+
+    event.sub_label = new_sub_label
+    event.save()
+    return make_response(
+        jsonify({"success": True, "message": "Event " + id + " sub label set to " + new_sub_label}), 200
+    )
 
 @bp.route("/events/<id>", methods=("DELETE",))
 def delete_event(id):
@@ -160,7 +186,7 @@ def delete_event(id):
         event = Event.get(Event.id == id)
     except DoesNotExist:
         return make_response(
-            jsonify({"success": False, "message": "Event" + id + " not found"}), 404
+            jsonify({"success": False, "message": "Event " + id + " not found"}), 404
         )
 
     media_name = f"{event.camera}-{event.id}"
@@ -175,7 +201,7 @@ def delete_event(id):
 
     event.delete_instance()
     return make_response(
-        jsonify({"success": True, "message": "Event" + id + " deleted"}), 200
+        jsonify({"success": True, "message": "Event " + id + " deleted"}), 200
     )
 
 
diff --git a/frigate/models.py b/frigate/models.py
index ff83dcb54..6c988d541 100644
--- a/frigate/models.py
+++ b/frigate/models.py
@@ -6,6 +6,7 @@ from playhouse.sqlite_ext import *
 class Event(Model):
     id = CharField(null=False, primary_key=True, max_length=30)
     label = CharField(index=True, max_length=20)
+    sub_label = CharField(max_length=20, null=True)
     camera = CharField(index=True, max_length=20)
     start_time = DateTimeField()
     end_time = DateTimeField()
diff --git a/migrations/008_add_sub_label.py b/migrations/008_add_sub_label.py
new file mode 100644
index 000000000..00603c63e
--- /dev/null
+++ b/migrations/008_add_sub_label.py
@@ -0,0 +1,46 @@
+"""Peewee migrations -- 008_add_sub_label.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,
+        sub_label=pw.CharField(max_length=20, null=True),
+    )
+
+
+def rollback(migrator, database, fake=False, **kwargs):
+    migrator.remove_fields(Event, ["sub_label"])
diff --git a/web/src/routes/Events.jsx b/web/src/routes/Events.jsx
index 02d6a5822..94a96debb 100644
--- a/web/src/routes/Events.jsx
+++ b/web/src/routes/Events.jsx
@@ -315,7 +315,7 @@ export default function Events({ path, ...props }) {
                     <div className="m-2 flex grow">
                       <div className="flex flex-col grow">
                         <div className="capitalize text-lg font-bold">
-                          {event.label} ({(event.top_score * 100).toFixed(0)}%)
+                          {event.sub_label ? `${event.label}: ${event.sub_label}` : event.label} ({(event.top_score * 100).toFixed(0)}%)
                         </div>
                         <div className="text-sm">
                           {new Date(event.start_time * 1000).toLocaleDateString()}{' '}