浏览代码

allow setting specific edgetpu in config

Blake Blackshear 4 年之前
父节点
当前提交
2f758af097
共有 3 个文件被更改,包括 22 次插入8 次删除
  1. 6 0
      config/config.example.yml
  2. 2 1
      detect_objects.py
  3. 14 7
      frigate/edgetpu.py

+ 6 - 0
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

+ 2 - 1
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 = {}

+ 14 - 7
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()