From 2f758af097f80e1a37784a4037e6df74fa352017 Mon Sep 17 00:00:00 2001 From: Blake Blackshear Date: Mon, 7 Sep 2020 12:49:47 -0500 Subject: [PATCH] allow setting specific edgetpu in config --- config/config.example.yml | 6 ++++++ detect_objects.py | 3 ++- frigate/edgetpu.py | 21 ++++++++++++++------- 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/config/config.example.yml b/config/config.example.yml index 180418e19..33e7af0bb 100644 --- a/config/config.example.yml +++ b/config/config.example.yml @@ -1,5 +1,11 @@ web_port: 5000 +################ +## Tell frigate to look for a specific EdgeTPU device. Useful if you want to run multiple instances of frigate +## on the same machine with multiple EdgeTPUs. https://coral.ai/docs/edgetpu/multiple-edgetpu/#using-the-tensorflow-lite-python-api +################ +tensorflow_device: usb + mqtt: host: mqtt.server.com topic_prefix: frigate diff --git a/detect_objects.py b/detect_objects.py index ab0bcc238..931ef01fd 100644 --- a/detect_objects.py +++ b/detect_objects.py @@ -61,6 +61,7 @@ GLOBAL_OBJECT_CONFIG = CONFIG.get('objects', {}) WEB_PORT = CONFIG.get('web_port', 5000) DEBUG = (CONFIG.get('debug', '0') == '1') +TENSORFLOW_DEVICE = CONFIG.get('tensorflow_device') def start_plasma_store(): plasma_cmd = ['plasma_store', '-m', '400000000', '-s', '/tmp/plasma'] @@ -190,7 +191,7 @@ def main(): event_queue = mp.Queue() # Start the shared tflite process - tflite_process = EdgeTPUProcess() + tflite_process = EdgeTPUProcess(TENSORFLOW_DEVICE) # start the camera processes camera_processes = {} diff --git a/frigate/edgetpu.py b/frigate/edgetpu.py index cffdb78aa..336746444 100644 --- a/frigate/edgetpu.py +++ b/frigate/edgetpu.py @@ -34,18 +34,24 @@ class ObjectDetector(ABC): pass class LocalObjectDetector(ObjectDetector): - def __init__(self, labels=None): + def __init__(self, tf_device=None, labels=None): if labels is None: self.labels = {} else: self.labels = load_labels(labels) + device_config = {"device": "usb"} + if not tf_device is None: + device_config = {"device": tf_device} + edge_tpu_delegate = None try: - edge_tpu_delegate = load_delegate('libedgetpu.so.1.0', {"device": "usb"}) - print("USB TPU found") + print(f"Attempting to load TPU as {device_config['device']}") + edge_tpu_delegate = load_delegate('libedgetpu.so.1.0', device_config) + print("TPU found") except ValueError: try: + print(f"Attempting to load TPU as pci:0") edge_tpu_delegate = load_delegate('libedgetpu.so.1.0', {"device": "pci:0"}) print("PCIe TPU found") except ValueError: @@ -92,11 +98,11 @@ class LocalObjectDetector(ObjectDetector): return detections -def run_detector(detection_queue, avg_speed, start): +def run_detector(detection_queue, avg_speed, start, tf_device): print(f"Starting detection process: {os.getpid()}") listen() plasma_client = plasma.connect("/tmp/plasma") - object_detector = LocalObjectDetector() + object_detector = LocalObjectDetector(tf_device=tf_device) while True: object_id_str = detection_queue.get() @@ -117,11 +123,12 @@ def run_detector(detection_queue, avg_speed, start): avg_speed.value = (avg_speed.value*9 + duration)/10 class EdgeTPUProcess(): - def __init__(self): + def __init__(self, tf_device=None): self.detection_queue = mp.Queue() self.avg_inference_speed = mp.Value('d', 0.01) self.detection_start = mp.Value('d', 0.0) self.detect_process = None + self.tf_device = tf_device self.start_or_restart() def start_or_restart(self): @@ -134,7 +141,7 @@ class EdgeTPUProcess(): print("Detection process didnt exit. Force killing...") self.detect_process.kill() self.detect_process.join() - self.detect_process = mp.Process(target=run_detector, args=(self.detection_queue, self.avg_inference_speed, self.detection_start)) + self.detect_process = mp.Process(target=run_detector, args=(self.detection_queue, self.avg_inference_speed, self.detection_start, self.tf_device)) self.detect_process.daemon = True self.detect_process.start()