administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-08-02 09:02:15 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Use sqlite-vec extension instead of chromadb for embeddings (#14163)

Browse Source 
* swap sqlite_vec for chroma in requirements

* load sqlite_vec in embeddings manager

* remove chroma and revamp Embeddings class for sqlite_vec

* manual minilm onnx inference

* remove chroma in clip model

* migrate api from chroma to sqlite_vec

* migrate event cleanup from chroma to sqlite_vec

* migrate embedding maintainer from chroma to sqlite_vec

* genai description for sqlite_vec

* load sqlite_vec in main thread db

* extend the SqliteQueueDatabase class and use peewee db.execute_sql

* search with Event type for similarity

* fix similarity search

* install and add comment about transformers

* fix normalization

* add id filter

* clean up

* clean up

* fully remove chroma and add transformers env var

* readd uvicorn for fastapi

* readd tokenizer parallelism env var

* remove chroma from docs

* remove chroma from UI

* try removing custom pysqlite3 build

* hard code limit

* optimize queries

* revert explore query

* fix query

* keep building pysqlite3

* single pass fetch and process

* remove unnecessary re-embed

* update deps

* move SqliteVecQueueDatabase to db directory

* make search thumbnail take up full size of results box

* improve typing

* improve model downloading and add status screen

* daemon downloading thread

* catch case when semantic search is disabled

* fix typing

* build sqlite_vec from source

* resolve conflict

* file permissions

* try build deps

* remove sources

* sources

* fix thread start

* include git in build

* reorder embeddings after detectors are started

* build with sqlite amalgamation

* non-platform specific

* use wget instead of curl

* remove unzip -d

* remove sqlite_vec from requirements and load the compiled version

* fix build

* avoid race in db connection

* add scale_factor and bias to description zscore normalization
This commit is contained in:
Josh Hawkins
2024-10-07 14:30:45 -06:00
committed by GitHub
parent 757150dec1
commit 24ac9f3e5a
42 changed files with 951 additions and 533 deletions
Download Patch File Download Diff File Expand all files Collapse all files
frigate/api/event.py
+85 -117
View File
@@ -1,8 +1,6 @@
"""Event apis."""
import base64
import datetime
import io
import logging
import os
from functools import reduce
@@ -10,12 +8,10 @@ from pathlib import Path
from urllib.parse import unquote
import cv2
import numpy as np
from fastapi import APIRouter, Request
from fastapi.params import Depends
from fastapi.responses import JSONResponse
from peewee import JOIN, DoesNotExist, fn, operator
from PIL import Image
from playhouse.shortcuts import model_to_dict
from frigate.api.defs.events_body import (
@@ -39,7 +35,6 @@ from frigate.const import (
CLIPS_DIR,
)
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
@@ -411,16 +406,12 @@ def events_search(request: Request, params: EventsSearchQueryParams = Depends())
event_filters = []
if cameras != "all":
camera_list = cameras.split(",")
event_filters.append((Event.camera << camera_list))
event_filters.append((Event.camera << cameras.split(",")))
if labels != "all":
label_list = labels.split(",")
event_filters.append((Event.label << label_list))
event_filters.append((Event.label << labels.split(",")))
if zones != "all":
# use matching so events with multiple zones
# still match on a search where any zone matches
zone_clauses = []
filtered_zones = zones.split(",")
@@ -431,8 +422,7 @@ def events_search(request: Request, params: EventsSearchQueryParams = Depends())
for zone in filtered_zones:
zone_clauses.append((Event.zones.cast("text") % f'*"{zone}"*'))
zone_clause = reduce(operator.or_, zone_clauses)
event_filters.append((zone_clause))
event_filters.append((reduce(operator.or_, zone_clauses)))
if after:
event_filters.append((Event.start_time > after))
@@ -441,13 +431,11 @@ def events_search(request: Request, params: EventsSearchQueryParams = Depends())
event_filters.append((Event.start_time < before))
if time_range != DEFAULT_TIME_RANGE:
# get timezone arg to ensure browser times are used
tz_name = params.timezone
hour_modifier, minute_modifier, _ = get_tz_modifiers(tz_name)
times = time_range.split(",")
time_after = times[0]
time_before = times[1]
time_after, time_before = times
start_hour_fun = fn.strftime(
"%H:%M",
@@ -470,132 +458,113 @@ def events_search(request: Request, params: EventsSearchQueryParams = Depends())
event_filters.append((start_hour_fun > time_after))
event_filters.append((start_hour_fun < time_before))
if event_filters:
filtered_event_ids = (
Event.select(Event.id)
.where(reduce(operator.and_, event_filters))
.tuples()
.iterator()
)
event_ids = [event_id[0] for event_id in filtered_event_ids]
if not event_ids:
return JSONResponse(content=[]) # No events to search on
else:
event_ids = []
# Build the Chroma where clause based on the event IDs
where = {"id": {"$in": event_ids}} if event_ids else {}
thumb_ids = {}
desc_ids = {}
# Perform semantic search
search_results = {}
if search_type == "similarity":
# Grab the ids of events that match the thumbnail image embeddings
try:
search_event: Event = Event.get(Event.id == event_id)
except DoesNotExist:
return JSONResponse(
content=(
{
"success": False,
"message": "Event not found",
}
),
content={
"success": False,
"message": "Event not found",
},
status_code=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_result = context.embeddings.search_thumbnail(search_event)
thumb_ids = dict(
zip(
thumb_result["ids"][0],
context.thumb_stats.normalize(thumb_result["distances"][0]),
[result[0] for result in thumb_result],
context.thumb_stats.normalize([result[1] for result in thumb_result]),
)
)
search_results = {
event_id: {"distance": distance, "source": "thumbnail"}
for event_id, distance in thumb_ids.items()
}
else:
search_types = search_type.split(",")
if "thumbnail" in search_types:
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_result = context.embeddings.search_thumbnail(query)
thumb_ids = dict(
zip(
thumb_result["ids"][0],
context.thumb_stats.normalize(thumb_result["distances"][0]),
[result[0] for result in thumb_result],
context.thumb_stats.normalize(
[result[1] for result in thumb_result]
),
)
)
search_results.update(
{
event_id: {"distance": distance, "source": "thumbnail"}
for event_id, distance in thumb_ids.items()
}
)
if "description" in search_types:
desc_result = context.embeddings.description.query(
query_texts=[query],
n_results=limit,
where=where,
)
desc_result = context.embeddings.search_description(query)
desc_ids = dict(
zip(
desc_result["ids"][0],
context.desc_stats.normalize(desc_result["distances"][0]),
[result[0] for result in desc_result],
context.desc_stats.normalize([result[1] for result in desc_result]),
)
)
for event_id, distance in desc_ids.items():
if (
event_id not in search_results
or distance < search_results[event_id]["distance"]
):
search_results[event_id] = {
"distance": distance,
"source": "description",
}
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:
if not search_results:
return JSONResponse(content=[])
# 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()
# Fetch events in a single query
events_query = Event.select(*selected_columns).join(
ReviewSegment,
JOIN.LEFT_OUTER,
on=(fn.json_extract(ReviewSegment.data, "$.detections").contains(Event.id)),
)
events = list(events)
events = [
{k: v for k, v in event.items() if k != "data"}
| {
"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]
# Apply filters, if any
if event_filters:
events_query = events_query.where(reduce(operator.and_, event_filters))
return JSONResponse(content=events)
# If we did a similarity search, limit events to those in search_results
if search_results:
events_query = events_query.where(Event.id << list(search_results.keys()))
# Fetch events and process them in a single pass
processed_events = []
for event in events_query.dicts():
processed_event = {k: v for k, v in event.items() if k != "data"}
processed_event["data"] = {
k: v
for k, v in event["data"].items()
if k in ["type", "score", "top_score", "description"]
}
if event["id"] in search_results:
processed_event["search_distance"] = search_results[event["id"]]["distance"]
processed_event["search_source"] = search_results[event["id"]]["source"]
processed_events.append(processed_event)
# Sort by search distance if search_results are available, otherwise by start_time
if search_results:
processed_events.sort(key=lambda x: x.get("search_distance", float("inf")))
else:
processed_events.sort(key=lambda x: x["start_time"], reverse=True)
# Limit the number of events returned
processed_events = processed_events[:limit]
return JSONResponse(content=processed_events)
@router.get("/events/summary")
@@ -975,10 +944,9 @@ def set_description(
# If semantic search is enabled, update the index
if request.app.frigate_config.semantic_search.enabled:
context: EmbeddingsContext = request.app.embeddings
context.embeddings.description.upsert(
documents=[new_description],
metadatas=[get_metadata(event)],
ids=[event_id],
context.embeddings.upsert_description(
event_id=event_id,
description=new_description,
)
response_message = (
@@ -1065,8 +1033,8 @@ def delete_event(request: Request, event_id: str):
# If semantic search is enabled, update the index
if request.app.frigate_config.semantic_search.enabled:
context: EmbeddingsContext = request.app.embeddings
context.embeddings.thumbnail.delete(ids=[event_id])
context.embeddings.description.delete(ids=[event_id])
context.embeddings.delete_thumbnail(id=[event_id])
context.embeddings.delete_description(id=[event_id])
return JSONResponse(
content=({"success": True, "message": "Event " + event_id + " deleted"}),
status_code=200,