config for model size

This commit is contained in:
Josh Hawkins 2025-04-15 09:32:51 -05:00
parent 1d54ae7bb2
commit 9476ca0dc3
2 changed files with 8 additions and 2 deletions

View File

@ -103,6 +103,9 @@ class LicensePlateRecognitionConfig(FrigateBaseModel):
default=LPRDeviceEnum.CPU, default=LPRDeviceEnum.CPU,
title="The device used for license plate recognition.", title="The device used for license plate recognition.",
) )
model_size: str = Field(
default="small", title="The size of the embeddings model used."
)
detection_threshold: float = Field( detection_threshold: float = Field(
default=0.7, default=0.7,
title="License plate object confidence score required to begin running recognition.", title="License plate object confidence score required to begin running recognition.",

View File

@ -31,11 +31,14 @@ class PaddleOCRDetection(BaseEmbedding):
requestor: InterProcessRequestor, requestor: InterProcessRequestor,
device: str = "AUTO", device: str = "AUTO",
): ):
model_file = (
"detection-large.onnx" if model_size == "large" else "detection-small.onnx"
)
super().__init__( super().__init__(
model_name="paddleocr-onnx", model_name="paddleocr-onnx",
model_file="detection.onnx", model_file=model_file,
download_urls={ download_urls={
"detection.onnx": "https://github.com/hawkeye217/paddleocr-onnx/raw/refs/heads/master/models/detection.onnx" model_file: f"https://github.com/hawkeye217/paddleocr-onnx/raw/refs/heads/master/models/{model_file}"
}, },
) )
self.requestor = requestor self.requestor = requestor