Semantic Search API (#12105)

* initial event search api implementation

* fix lint

* fix tests

* move chromadb imports and pysqlite hotswap to fix tests

* remove unused import

* switch default limit to 50

* fix events accidently pulling inside chroma results loop

frigate/__main__.py

@@ -1,12 +1,9 @@
 import faulthandler
-import sys
 import threading
 
 from flask import cli
 
-# Hotsawp the sqlite3 module for Chroma compatibility
-__import__("pysqlite3")
-sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
+from frigate.app import FrigateApp
 
 faulthandler.enable()
 
@@ -15,8 +12,6 @@
 cli.show_server_banner = lambda *x: None
 
 if __name__ == "__main__":
-    from frigate.app import FrigateApp
-
     frigate_app = FrigateApp()
 
     frigate_app.start()

frigate/api/app.py

@@ -23,6 +23,7 @@
 from frigate.api.review import ReviewBp
 from frigate.config import FrigateConfig
 from frigate.const import CONFIG_DIR
+from frigate.embeddings import EmbeddingsContext
 from frigate.events.external import ExternalEventProcessor
 from frigate.models import Event, Timeline
 from frigate.plus import PlusApi
@@ -52,6 +53,7 @@
 def create_app(
     frigate_config,
     database: SqliteQueueDatabase,
+    embeddings: EmbeddingsContext,
     detected_frames_processor,
     storage_maintainer: StorageMaintainer,
     onvif: OnvifController,
@@ -79,6 +81,7 @@ def _db_close(exc):
             database.close()
 
     app.frigate_config = frigate_config
+    app.embeddings = embeddings
     app.detected_frames_processor = detected_frames_processor
     app.storage_maintainer = storage_maintainer
     app.onvif = onvif

frigate/api/event.py

@@ -1,5 +1,7 @@
 """Event apis."""
 
+import base64
+import io
 import logging
 import os
 from datetime import datetime
@@ -8,20 +10,24 @@
 from urllib.parse import unquote
 
 import cv2
+import numpy as np
 from flask import (
     Blueprint,
     current_app,
     jsonify,
     make_response,
     request,
 )
-from peewee import DoesNotExist, fn, operator
+from peewee import JOIN, DoesNotExist, fn, operator
+from PIL import Image
 from playhouse.shortcuts import model_to_dict
 
 from frigate.const import (
     CLIPS_DIR,
 )
-from frigate.models import Event, Timeline
+from frigate.embeddings import EmbeddingsContext
+from frigate.embeddings.embeddings import get_metadata
+from frigate.models import Event, ReviewSegment, Timeline
 from frigate.object_processing import TrackedObject
 from frigate.util.builtin import get_tz_modifiers
 
@@ -245,6 +251,189 @@ def events():
     return jsonify(list(events))
 
 
+@EventBp.route("/events/search")
+def events_search():
+    query = request.args.get("query", type=str)
+    search_type = request.args.get("search_type", "text", type=str)
+    include_thumbnails = request.args.get("include_thumbnails", default=1, type=int)
+    limit = request.args.get("limit", 50, type=int)
+
+    # Filters
+    cameras = request.args.get("cameras", "all", type=str)
+    labels = request.args.get("labels", "all", type=str)
+    zones = request.args.get("zones", "all", type=str)
+    after = request.args.get("after", type=float)
+    before = request.args.get("before", type=float)
+
+    if not query:
+        return make_response(
+            jsonify(
+                {
+                    "success": False,
+                    "message": "A search query must be supplied",
+                }
+            ),
+            400,
+        )
+
+    if not current_app.frigate_config.semantic_search.enabled:
+        return make_response(
+            jsonify(
+                {
+                    "success": False,
+                    "message": "Semantic search is not enabled",
+                }
+            ),
+            400,
+        )
+
+    context: EmbeddingsContext = current_app.embeddings
+
+    selected_columns = [
+        Event.id,
+        Event.camera,
+        Event.label,
+        Event.sub_label,
+        Event.zones,
+        Event.start_time,
+        Event.end_time,
+        Event.data,
+        ReviewSegment.thumb_path,
+    ]
+
+    if include_thumbnails:
+        selected_columns.append(Event.thumbnail)
+
+    # Build the where clause for the embeddings query
+    embeddings_filters = []
+
+    if cameras != "all":
+        camera_list = cameras.split(",")
+        embeddings_filters.append({"camera": {"$in": camera_list}})
+
+    if labels != "all":
+        label_list = labels.split(",")
+        embeddings_filters.append({"label": {"$in": label_list}})
+
+    if zones != "all":
+        filtered_zones = zones.split(",")
+        zone_filters = [{f"zones_{zone}": {"$eq": True}} for zone in filtered_zones]
+        if len(zone_filters) > 1:
+            embeddings_filters.append({"$or": zone_filters})
+        else:
+            embeddings_filters.append(zone_filters[0])
+
+    if after:
+        embeddings_filters.append({"start_time": {"$gt": after}})
+
+    if before:
+        embeddings_filters.append({"start_time": {"$lt": before}})
+
+    where = None
+    if len(embeddings_filters) > 1:
+        where = {"$and": embeddings_filters}
+    elif len(embeddings_filters) == 1:
+        where = embeddings_filters[0]
+
+    thumb_ids = {}
+    desc_ids = {}
+
+    if search_type == "thumbnail":
+        # Grab the ids of events that match the thumbnail image embeddings
+        try:
+            search_event: Event = Event.get(Event.id == query)
+        except DoesNotExist:
+            return make_response(
+                jsonify(
+                    {
+                        "success": False,
+                        "message": "Event not found",
+                    }
+                ),
+                404,
+            )
+        thumbnail = base64.b64decode(search_event.thumbnail)
+        img = np.array(Image.open(io.BytesIO(thumbnail)).convert("RGB"))
+        thumb_result = context.embeddings.thumbnail.query(
+            query_images=[img],
+            n_results=limit,
+            where=where,
+        )
+        thumb_ids = dict(zip(thumb_result["ids"][0], thumb_result["distances"][0]))
+    else:
+        thumb_result = context.embeddings.thumbnail.query(
+            query_texts=[query],
+            n_results=limit,
+            where=where,
+        )
+        # Do a rudimentary normalization of the difference in distances returned by CLIP and MiniLM.
+        thumb_ids = dict(
+            zip(
+                thumb_result["ids"][0],
+                context.thumb_stats.normalize(thumb_result["distances"][0]),
+            )
+        )
+        desc_result = context.embeddings.description.query(
+            query_texts=[query],
+            n_results=limit,
+            where=where,
+        )
+        desc_ids = dict(
+            zip(
+                desc_result["ids"][0],
+                context.desc_stats.normalize(desc_result["distances"][0]),
+            )
+        )
+
+    results = {}
+    for event_id in thumb_ids.keys() | desc_ids:
+        min_distance = min(
+            i
+            for i in (thumb_ids.get(event_id), desc_ids.get(event_id))
+            if i is not None
+        )
+        results[event_id] = {
+            "distance": min_distance,
+            "source": "thumbnail"
+            if min_distance == thumb_ids.get(event_id)
+            else "description",
+        }
+
+    if not results:
+        return jsonify([])
+
+    # Get the event data
+    events = (
+        Event.select(*selected_columns)
+        .join(
+            ReviewSegment,
+            JOIN.LEFT_OUTER,
+            on=(fn.json_extract(ReviewSegment.data, "$.detections").contains(Event.id)),
+        )
+        .where(Event.id << list(results.keys()))
+        .dicts()
+        .iterator()
+    )
+    events = list(events)
+
+    events = [
+        {k: v for k, v in event.items() if k != "data"}
+        | {
+            k: v
+            for k, v in event["data"].items()
+            if k in ["type", "score", "top_score", "description"]
+        }
+        | {
+            "search_distance": results[event["id"]]["distance"],
+            "search_source": results[event["id"]]["source"],
+        }
+        for event in events
+    ]
+    events = sorted(events, key=lambda x: x["search_distance"])[:limit]
+
+    return jsonify(events)
+
+
 @EventBp.route("/events/summary")
 def events_summary():
     tz_name = request.args.get("timezone", default="utc", type=str)
@@ -604,6 +793,52 @@ def set_sub_label(id):
     )
 
 
+@EventBp.route("/events/<id>/description", methods=("POST",))
+def set_description(id):
+    try:
+        event: Event = Event.get(Event.id == id)
+    except DoesNotExist:
+        return make_response(
+            jsonify({"success": False, "message": "Event " + id + " not found"}), 404
+        )
+
+    json: dict[str, any] = request.get_json(silent=True) or {}
+    new_description = json.get("description")
+
+    if new_description is None or len(new_description) == 0:
+        return make_response(
+            jsonify(
+                {
+                    "success": False,
+                    "message": "description cannot be empty",
+                }
+            ),
+            400,
+        )
+
+    event.data["description"] = new_description
+    event.save()
+
+    # If semantic search is enabled, update the index
+    if current_app.frigate_config.semantic_search.enabled:
+        context: EmbeddingsContext = current_app.embeddings
+        context.embeddings.description.upsert(
+            documents=[new_description],
+            metadatas=[get_metadata(event)],
+            ids=[id],
+        )
+
+    return make_response(
+        jsonify(
+            {
+                "success": True,
+                "message": "Event " + id + " description set to " + new_description,
+            }
+        ),
+        200,
+    )
+
+
 @EventBp.route("/events/<id>", methods=("DELETE",))
 def delete_event(id):
     try:
@@ -625,6 +860,11 @@ def delete_event(id):
 
     event.delete_instance()
     Timeline.delete().where(Timeline.source_id == id).execute()
+    # If semantic search is enabled, update the index
+    if current_app.frigate_config.semantic_search.enabled:
+        context: EmbeddingsContext = current_app.embeddings
+        context.embeddings.thumbnail.delete(ids=[id])
+        context.embeddings.description.delete(ids=[id])
     return make_response(
         jsonify({"success": True, "message": "Event " + id + " deleted"}), 200
     )

frigate/app.py

@@ -37,8 +37,7 @@
     MODEL_CACHE_DIR,
     RECORD_DIR,
 )
-from frigate.embeddings import manage_embeddings
-from frigate.embeddings.embeddings import Embeddings
+from frigate.embeddings import EmbeddingsContext, manage_embeddings
 from frigate.events.audio import listen_to_audio
 from frigate.events.cleanup import EventCleanup
 from frigate.events.external import ExternalEventProcessor
@@ -322,7 +321,7 @@ def init_review_segment_manager(self) -> None:
 
     def init_embeddings_manager(self) -> None:
         # Create a client for other processes to use
-        self.embeddings = Embeddings()
+        self.embeddings = EmbeddingsContext()
         embedding_process = mp.Process(
             target=manage_embeddings,
             name="embeddings_manager",
@@ -384,6 +383,7 @@ def init_web_server(self) -> None:
         self.flask_app = create_app(
             self.config,
             self.db,
+            self.embeddings,
             self.detected_frames_processor,
             self.storage_maintainer,
             self.onvif_controller,
@@ -811,6 +811,9 @@ def stop(self) -> None:
         self.frigate_watchdog.join()
         self.db.stop()
 
+        # Save embeddings stats to disk
+        self.embeddings.save_stats()
+
         # Stop Communicators
         self.inter_process_communicator.stop()
         self.inter_config_updater.stop()