mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-09-28 19:06:52 +03:00
121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
import logging
|
|
import os.path
|
|
import re
|
|
import urllib.request
|
|
from typing import Literal
|
|
|
|
from pydantic import ConfigDict
|
|
|
|
from frigate.const import MODEL_CACHE_DIR
|
|
from frigate.detectors.detection_api import DetectionApi
|
|
from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum
|
|
from frigate.util.model import post_process_yolo
|
|
from frigate.util.runtime_deps import Artifact, ArtifactKind, RuntimeManifest
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DETECTOR_KEY = "axengine"
|
|
|
|
# The AXEngine python package is installed at first start rather than shipped
|
|
# in the image; its native libraries are bind mounted from the host.
|
|
AXENGINE_VERSION = "0.1.3"
|
|
AXENGINE_MANIFEST = RuntimeManifest(
|
|
name=DETECTOR_KEY,
|
|
version=AXENGINE_VERSION,
|
|
artifacts=(
|
|
Artifact(
|
|
url=f"https://github.com/AXERA-TECH/pyaxengine/releases/download/{AXENGINE_VERSION}-frigate/axengine-{AXENGINE_VERSION}-py3-none-any.whl",
|
|
sha256="e995b8a887b067dc3456512aae2fa9c84f70e708c28b11caf184efdc254c64ae",
|
|
kind=ArtifactKind.wheel,
|
|
),
|
|
),
|
|
import_check="axengine",
|
|
)
|
|
|
|
supported_models = {
|
|
ModelTypeEnum.yologeneric: "frigate-yolov9-.*$",
|
|
}
|
|
|
|
model_cache_dir = os.path.join(MODEL_CACHE_DIR, "axengine_cache/")
|
|
|
|
|
|
class AxengineDetectorConfig(BaseDetectorConfig):
|
|
"""AXERA AX650N/AX8850N NPU detector running compiled .axmodel files via the AXEngine runtime."""
|
|
|
|
model_config = ConfigDict(
|
|
title="AXEngine NPU",
|
|
)
|
|
|
|
type: Literal[DETECTOR_KEY]
|
|
|
|
|
|
class Axengine(DetectionApi):
|
|
type_key = DETECTOR_KEY
|
|
runtime_manifest = AXENGINE_MANIFEST
|
|
|
|
def __init__(self, config: AxengineDetectorConfig):
|
|
self.activate_dependencies()
|
|
|
|
try:
|
|
import axengine as axe
|
|
except ModuleNotFoundError:
|
|
raise ImportError(
|
|
"AXEngine is not installed. Frigate installs it at startup when an "
|
|
"axengine detector is configured; check the startup log for errors."
|
|
) from None
|
|
|
|
logger.info("__init__ axengine")
|
|
super().__init__(config)
|
|
self.height = config.model.height
|
|
self.width = config.model.width
|
|
model_path = config.model.path or "frigate-yolov9-tiny"
|
|
model_props = self.parse_model_input(model_path)
|
|
self.session = axe.InferenceSession(model_props["path"])
|
|
|
|
def __del__(self):
|
|
pass
|
|
|
|
def parse_model_input(self, model_path):
|
|
model_props = {}
|
|
model_props["preset"] = True
|
|
|
|
model_matched = False
|
|
|
|
for model_type, pattern in supported_models.items():
|
|
if re.match(pattern, model_path):
|
|
model_matched = True
|
|
model_props["model_type"] = model_type
|
|
|
|
if model_matched:
|
|
model_props["filename"] = model_path + ".axmodel"
|
|
model_props["path"] = model_cache_dir + model_props["filename"]
|
|
|
|
if not os.path.isfile(model_props["path"]):
|
|
self.download_model(model_props["filename"])
|
|
else:
|
|
supported_models_str = ", ".join(model[1:-1] for model in supported_models)
|
|
raise Exception(
|
|
f"Model {model_path} is unsupported. Provide your own model or choose one of the following: {supported_models_str}"
|
|
)
|
|
return model_props
|
|
|
|
def download_model(self, filename):
|
|
if not os.path.isdir(model_cache_dir):
|
|
os.mkdir(model_cache_dir)
|
|
|
|
HF_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co")
|
|
urllib.request.urlretrieve(
|
|
f"{HF_ENDPOINT}/AXERA-TECH/frigate-resource/resolve/axmodel/{filename}",
|
|
model_cache_dir + filename,
|
|
)
|
|
|
|
def detect_raw(self, tensor_input):
|
|
results = None
|
|
results = self.session.run(None, {"images": tensor_input})
|
|
if self.detector_config.model.model_type == ModelTypeEnum.yologeneric:
|
|
return post_process_yolo(results, self.width, self.height)
|
|
else:
|
|
raise ValueError(
|
|
f'Model type "{self.detector_config.model.model_type}" is currently not supported.'
|
|
)
|