forked from Github/frigate
Allow using full tensorflow in cpu/edgetpu detector plugins (#5611)
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.
This commit is contained in:
@@ -5,7 +5,11 @@ from frigate.detectors.detection_api import DetectionApi
|
||||
from frigate.detectors.detector_config import BaseDetectorConfig
|
||||
from typing import Literal
|
||||
from pydantic import Extra, Field
|
||||
import tflite_runtime.interpreter as tflite
|
||||
|
||||
try:
|
||||
from tflite_runtime.interpreter import Interpreter
|
||||
except ModuleNotFoundError:
|
||||
from tensorflow.lite.python.interpreter import Interpreter
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -22,7 +26,7 @@ class CpuTfl(DetectionApi):
|
||||
type_key = DETECTOR_KEY
|
||||
|
||||
def __init__(self, detector_config: CpuDetectorConfig):
|
||||
self.interpreter = tflite.Interpreter(
|
||||
self.interpreter = Interpreter(
|
||||
model_path=detector_config.model.path or "/cpu_model.tflite",
|
||||
num_threads=detector_config.num_threads or 3,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user