From 2381f7a754070280a16d84fede71343b451f0be5 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Mon, 21 Oct 2024 15:29:26 -0600
Subject: [PATCH] Cleanup preprocessing

---
 frigate/embeddings/embeddings.py     | 9 ++-------
 frigate/embeddings/functions/onnx.py | 2 ++
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/frigate/embeddings/embeddings.py b/frigate/embeddings/embeddings.py
index 9ee508823..8707f6f37 100644
--- a/frigate/embeddings/embeddings.py
+++ b/frigate/embeddings/embeddings.py
@@ -125,8 +125,7 @@ class Embeddings:
 
     def upsert_thumbnail(self, event_id: str, thumbnail: bytes) -> ndarray:
         # Convert thumbnail bytes to PIL Image
-        image = Image.open(io.BytesIO(thumbnail)).convert("RGB")
-        embedding = self.vision_embedding([image])[0]
+        embedding = self.vision_embedding([thumbnail])[0]
 
         self.db.execute_sql(
             """
@@ -139,12 +138,8 @@ class Embeddings:
         return embedding
 
     def batch_upsert_thumbnail(self, event_thumbs: dict[str, bytes]) -> list[ndarray]:
-        images = [
-            Image.open(io.BytesIO(thumb)).convert("RGB")
-            for thumb in event_thumbs.values()
-        ]
         ids = list(event_thumbs.keys())
-        embeddings = self.vision_embedding(images)
+        embeddings = self.vision_embedding(list(event_thumbs.values()))
 
         items = []
 
diff --git a/frigate/embeddings/functions/onnx.py b/frigate/embeddings/functions/onnx.py
index c669bcf73..b79f2fceb 100644
--- a/frigate/embeddings/functions/onnx.py
+++ b/frigate/embeddings/functions/onnx.py
@@ -170,6 +170,8 @@ class GenericONNXEmbedding:
             if image.startswith("http"):
                 response = requests.get(image)
                 image = Image.open(BytesIO(response.content)).convert("RGB")
+        elif isinstance(image, bytes):
+            image = Image.open(BytesIO(image)).convert("RGB")
 
         return image