Refactor ONNX embedding class to use a base class and type-specific c…

…lasses (#16703) * Move onnx runner * Build out base embedding * Convert text embedding to separate class * Move image embedding to separate * Move LPR to separate class * Remove mono embedding * Simplify model downloading * Reorganize jina v1 embeddings * Cleanup * Cleanup for review
blakeblackshear · Feb 20, 2025 · c736b1d · c736b1d
1 parent 649e5cf
commit c736b1d
Show file tree

Hide file tree

Showing 8 changed files with 701 additions and 449 deletions.
diff --git a/frigate/data_processing/real_time/license_plate_processor.py b/frigate/data_processing/real_time/license_plate_processor.py
@@ -16,7 +16,12 @@
 from frigate.comms.inter_process import InterProcessRequestor
 from frigate.config import FrigateConfig
 from frigate.const import FRIGATE_LOCALHOST
-from frigate.embeddings.functions.onnx import GenericONNXEmbedding, ModelTypeEnum
+from frigate.embeddings.onnx.lpr_embedding import (
+    LicensePlateDetector,
+    PaddleOCRClassification,
+    PaddleOCRDetection,
+    PaddleOCRRecognition,
+)
 from frigate.util.image import area
 
 from ..types import DataProcessorMetrics
@@ -52,49 +57,26 @@ def __init__(self, config: FrigateConfig, metrics: DataProcessorMetrics):
         self.lpr_recognition_model = None
 
         if self.config.lpr.enabled:
-            self.detection_model = GenericONNXEmbedding(
-                model_name="paddleocr-onnx",
-                model_file="detection.onnx",
-                download_urls={
-                    "detection.onnx": "https://github.com/hawkeye217/paddleocr-onnx/raw/refs/heads/master/models/detection.onnx"
-                },
+            self.detection_model = PaddleOCRDetection(
                 model_size="large",
-                model_type=ModelTypeEnum.lpr_detect,
                 requestor=self.requestor,
                 device="CPU",
             )
 
-            self.classification_model = GenericONNXEmbedding(
-                model_name="paddleocr-onnx",
-                model_file="classification.onnx",
-                download_urls={
-                    "classification.onnx": "https://github.com/hawkeye217/paddleocr-onnx/raw/refs/heads/master/models/classification.onnx"
-                },
+            self.classification_model = PaddleOCRClassification(
                 model_size="large",
-                model_type=ModelTypeEnum.lpr_classify,
                 requestor=self.requestor,
                 device="CPU",
             )
 
-            self.recognition_model = GenericONNXEmbedding(
-                model_name="paddleocr-onnx",
-                model_file="recognition.onnx",
-                download_urls={
-                    "recognition.onnx": "https://github.com/hawkeye217/paddleocr-onnx/raw/refs/heads/master/models/recognition.onnx"
-                },
+            self.recognition_model = PaddleOCRRecognition(
                 model_size="large",
-                model_type=ModelTypeEnum.lpr_recognize,
                 requestor=self.requestor,
                 device="CPU",
             )
-            self.yolov9_detection_model = GenericONNXEmbedding(
-                model_name="yolov9_license_plate",
-                model_file="yolov9-256-license-plates.onnx",
-                download_urls={
-                    "yolov9-256-license-plates.onnx": "https://github.com/hawkeye217/yolov9-license-plates/raw/refs/heads/master/models/yolov9-256-license-plates.onnx"
-                },
+
+            self.yolov9_detection_model = LicensePlateDetector(
                 model_size="large",
-                model_type=ModelTypeEnum.yolov9_lpr_detect,
                 requestor=self.requestor,
                 device="CPU",
             )

diff --git a/frigate/embeddings/embeddings.py b/frigate/embeddings/embeddings.py
@@ -22,7 +22,7 @@
 from frigate.util.builtin import serialize
 from frigate.util.path import get_event_thumbnail_bytes
 
-from .functions.onnx import GenericONNXEmbedding, ModelTypeEnum
+from .onnx.jina_v1_embedding import JinaV1ImageEmbedding, JinaV1TextEmbedding
 
 logger = logging.getLogger(__name__)
 
@@ -97,36 +97,14 @@ def __init__(
                 },
             )
 
-        self.text_embedding = GenericONNXEmbedding(
-            model_name="jinaai/jina-clip-v1",
-            model_file="text_model_fp16.onnx",
-            tokenizer_file="tokenizer",
-            download_urls={
-                "text_model_fp16.onnx": "https://huggingface.co/jinaai/jina-clip-v1/resolve/main/onnx/text_model_fp16.onnx",
-            },
+        self.text_embedding = JinaV1TextEmbedding(
             model_size=config.semantic_search.model_size,
-            model_type=ModelTypeEnum.text,
             requestor=self.requestor,
             device="CPU",
         )
 
-        model_file = (
-            "vision_model_fp16.onnx"
-            if self.config.semantic_search.model_size == "large"
-            else "vision_model_quantized.onnx"
-        )
-
-        download_urls = {
-            model_file: f"https://huggingface.co/jinaai/jina-clip-v1/resolve/main/onnx/{model_file}",
-            "preprocessor_config.json": "https://huggingface.co/jinaai/jina-clip-v1/resolve/main/preprocessor_config.json",
-        }
-
-        self.vision_embedding = GenericONNXEmbedding(
-            model_name="jinaai/jina-clip-v1",
-            model_file=model_file,
-            download_urls=download_urls,
+        self.vision_embedding = JinaV1ImageEmbedding(
             model_size=config.semantic_search.model_size,
-            model_type=ModelTypeEnum.vision,
             requestor=self.requestor,
             device="GPU" if config.semantic_search.model_size == "large" else "CPU",
         )