mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-10-15 20:07:23 +02:00
cd5f4b1534
* detectors/edgetpu: add support for yolov8 models * docs: edgetpu yolov8 running * docs: edgetpu yolov8 attribution and language * Update docs/docs/configuration/object_detectors.md Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com> --------- Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
import logging
|
|
|
|
import numpy as np
|
|
from pydantic import Field
|
|
from typing_extensions import Literal
|
|
|
|
from frigate.detectors.detection_api import DetectionApi
|
|
from frigate.detectors.detector_config import BaseDetectorConfig
|
|
from frigate.detectors.util import yolov8_postprocess
|
|
|
|
try:
|
|
from tflite_runtime.interpreter import Interpreter, load_delegate
|
|
except ModuleNotFoundError:
|
|
from tensorflow.lite.python.interpreter import Interpreter, load_delegate
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DETECTOR_KEY = "edgetpu"
|
|
|
|
|
|
class EdgeTpuDetectorConfig(BaseDetectorConfig):
|
|
type: Literal[DETECTOR_KEY]
|
|
device: str = Field(default=None, title="Device Type")
|
|
|
|
|
|
class EdgeTpuTfl(DetectionApi):
|
|
type_key = DETECTOR_KEY
|
|
|
|
def __init__(self, detector_config: EdgeTpuDetectorConfig):
|
|
device_config = {}
|
|
if detector_config.device is not None:
|
|
device_config = {"device": detector_config.device}
|
|
|
|
edge_tpu_delegate = None
|
|
|
|
try:
|
|
device_type = (
|
|
device_config["device"] if "device" in device_config else "auto"
|
|
)
|
|
logger.info(f"Attempting to load TPU as {device_type}")
|
|
edge_tpu_delegate = load_delegate("libedgetpu.so.1.0", device_config)
|
|
logger.info("TPU found")
|
|
self.interpreter = Interpreter(
|
|
model_path=detector_config.model.path,
|
|
experimental_delegates=[edge_tpu_delegate],
|
|
)
|
|
except ValueError:
|
|
logger.error(
|
|
"No EdgeTPU was detected. If you do not have a Coral device yet, you must configure CPU detectors."
|
|
)
|
|
raise
|
|
|
|
self.interpreter.allocate_tensors()
|
|
|
|
self.tensor_input_details = self.interpreter.get_input_details()
|
|
self.tensor_output_details = self.interpreter.get_output_details()
|
|
self.model_type = detector_config.model.model_type
|
|
|
|
def detect_raw(self, tensor_input):
|
|
if self.model_type == "yolov8":
|
|
scale, zero_point = self.tensor_input_details[0]["quantization"]
|
|
tensor_input = (
|
|
(tensor_input - scale * zero_point * 255) * (1.0 / (scale * 255))
|
|
).astype(self.tensor_input_details[0]["dtype"])
|
|
|
|
self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
|
|
self.interpreter.invoke()
|
|
|
|
if self.model_type == "yolov8":
|
|
scale, zero_point = self.tensor_output_details[0]["quantization"]
|
|
tensor_output = self.interpreter.get_tensor(
|
|
self.tensor_output_details[0]["index"]
|
|
)
|
|
tensor_output = (tensor_output.astype(np.float32) - zero_point) * scale
|
|
model_input_shape = self.tensor_input_details[0]["shape"]
|
|
tensor_output[:, [0, 2]] *= model_input_shape[2]
|
|
tensor_output[:, [1, 3]] *= model_input_shape[1]
|
|
return yolov8_postprocess(model_input_shape, tensor_output)
|
|
|
|
boxes = self.interpreter.tensor(self.tensor_output_details[0]["index"])()[0]
|
|
class_ids = self.interpreter.tensor(self.tensor_output_details[1]["index"])()[0]
|
|
scores = self.interpreter.tensor(self.tensor_output_details[2]["index"])()[0]
|
|
count = int(
|
|
self.interpreter.tensor(self.tensor_output_details[3]["index"])()[0]
|
|
)
|
|
|
|
detections = np.zeros((20, 6), np.float32)
|
|
|
|
for i in range(count):
|
|
if scores[i] < 0.4 or i == 20:
|
|
break
|
|
detections[i] = [
|
|
class_ids[i],
|
|
float(scores[i]),
|
|
boxes[i][0],
|
|
boxes[i][1],
|
|
boxes[i][2],
|
|
boxes[i][3],
|
|
]
|
|
|
|
return detections
|