|
@@ -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()
|
|
|
|