mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-11-21 19:07:46 +01:00
47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
|
import logging
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from frigate.detectors.detection_api import DetectionApi
|
||
|
|
import tflite_runtime.interpreter as tflite
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class CpuTfl(DetectionApi):
|
||
|
|
def __init__(self, det_device=None, model_config=None, num_threads=3):
|
||
|
|
self.interpreter = tflite.Interpreter(
|
||
|
|
model_path=model_config.path or "/cpu_model.tflite", num_threads=num_threads
|
||
|
|
)
|
||
|
|
|
||
|
|
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
|