mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-07-27 06:09:01 +03:00
Jozef HuscavaandGitHub
6f24f5a595
* Convert face crops to RGB before embedding Face crops flow through the cv2 pipeline as BGR arrays, but _process_image passes ndarrays to PIL without any channel conversion, so the FaceNet and ArcFace embedders receive BGR input while both models expect RGB. The error is symmetric between enrollment and recognition so it partially cancels, but it still costs accuracy. * Move BGR to RGB conversion into a shared helper Deduplicate the channel swap from both _preprocess_inputs methods into a BaseEmbedding._bgr_to_rgb static helper, as suggested in review.
104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
"""Base class for onnx embedding implementations."""
|
|
|
|
import logging
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
from io import BytesIO
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import requests
|
|
from PIL import Image
|
|
|
|
from frigate.const import UPDATE_MODEL_STATE
|
|
from frigate.types import ModelStatusTypesEnum
|
|
from frigate.util.downloader import ModelDownloader
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BaseEmbedding(ABC):
|
|
"""Base embedding class."""
|
|
|
|
def __init__(self, model_name: str, model_file: str, download_urls: dict[str, str]):
|
|
self.model_name = model_name
|
|
self.model_file = model_file
|
|
self.download_urls = download_urls
|
|
self.downloader: ModelDownloader = None
|
|
|
|
def _download_model(self, path: str):
|
|
try:
|
|
file_name = os.path.basename(path)
|
|
|
|
if file_name in self.download_urls:
|
|
ModelDownloader.download_from_url(self.download_urls[file_name], path)
|
|
|
|
self.downloader.requestor.send_data(
|
|
UPDATE_MODEL_STATE,
|
|
{
|
|
"model": f"{self.model_name}-{file_name}",
|
|
"state": ModelStatusTypesEnum.downloaded,
|
|
},
|
|
)
|
|
except Exception:
|
|
self.downloader.requestor.send_data(
|
|
UPDATE_MODEL_STATE,
|
|
{
|
|
"model": f"{self.model_name}-{file_name}",
|
|
"state": ModelStatusTypesEnum.error,
|
|
},
|
|
)
|
|
|
|
@abstractmethod
|
|
def _load_model_and_utils(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def _preprocess_inputs(self, raw_inputs: Any) -> Any:
|
|
pass
|
|
|
|
@staticmethod
|
|
def _bgr_to_rgb(frame: Any) -> Any:
|
|
if isinstance(frame, np.ndarray) and frame.ndim == 3:
|
|
return np.ascontiguousarray(frame[:, :, ::-1])
|
|
return frame
|
|
|
|
def _process_image(self, image, output: str = "RGB") -> Image.Image:
|
|
if isinstance(image, str):
|
|
if image.startswith("http"):
|
|
response = requests.get(image)
|
|
image = Image.open(BytesIO(response.content)).convert(output)
|
|
elif isinstance(image, bytes):
|
|
image = Image.open(BytesIO(image)).convert(output)
|
|
elif isinstance(image, np.ndarray):
|
|
image = Image.fromarray(image)
|
|
|
|
return image
|
|
|
|
def _postprocess_outputs(self, outputs: Any) -> Any:
|
|
return outputs
|
|
|
|
def __call__(
|
|
self, inputs: list[str] | list[Image.Image] | list[str]
|
|
) -> list[np.ndarray]:
|
|
self._load_model_and_utils()
|
|
processed = self._preprocess_inputs(inputs)
|
|
input_names = self.runner.get_input_names()
|
|
onnx_inputs = {name: [] for name in input_names}
|
|
input: dict[str, Any]
|
|
for input in processed:
|
|
for key, value in input.items():
|
|
if key in input_names:
|
|
onnx_inputs[key].append(value[0])
|
|
|
|
for key in input_names:
|
|
if onnx_inputs.get(key):
|
|
onnx_inputs[key] = np.stack(onnx_inputs[key])
|
|
else:
|
|
logger.warning(f"Expected input '{key}' not found in onnx_inputs")
|
|
|
|
outputs = self.runner.run(onnx_inputs)[0]
|
|
embeddings = self._postprocess_outputs(outputs)
|
|
|
|
return [embedding for embedding in embeddings]
|