mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-11-21 19:07:46 +01:00
161e7b3fd7
It supports the same entrypoints, given that tflite is a small cut-out of the big tensorflow picture. This patch was created for downstream usage in nixpkgs, where we don't have the tflite python package, but do have the full tensorflow package.
65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
import logging
|
|
import numpy as np
|
|
|
|
from frigate.detectors.detection_api import DetectionApi
|
|
from frigate.detectors.detector_config import BaseDetectorConfig
|
|
from typing import Literal
|
|
from pydantic import Extra, Field
|
|
|
|
try:
|
|
from tflite_runtime.interpreter import Interpreter
|
|
except ModuleNotFoundError:
|
|
from tensorflow.lite.python.interpreter import Interpreter
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DETECTOR_KEY = "cpu"
|
|
|
|
|
|
class CpuDetectorConfig(BaseDetectorConfig):
|
|
type: Literal[DETECTOR_KEY]
|
|
num_threads: int = Field(default=3, title="Number of detection threads")
|
|
|
|
|
|
class CpuTfl(DetectionApi):
|
|
type_key = DETECTOR_KEY
|
|
|
|
def __init__(self, detector_config: CpuDetectorConfig):
|
|
self.interpreter = Interpreter(
|
|
model_path=detector_config.model.path or "/cpu_model.tflite",
|
|
num_threads=detector_config.num_threads or 3,
|
|
)
|
|
|
|
self.interpreter.allocate_tensors()
|
|
|
|
self.tensor_input_details = self.interpreter.get_input_details()
|
|
self.tensor_output_details = self.interpreter.get_output_details()
|
|
|
|
def detect_raw(self, tensor_input):
|
|
self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
|
|
self.interpreter.invoke()
|
|
|
|
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
|