support multiple coral devices (fixes #100)

This commit is contained in:
Blake Blackshear
2020-10-10 06:57:43 -05:00
parent 49fca1b839
commit f946813ccb
5 changed files with 91 additions and 75 deletions

View File

@@ -48,15 +48,12 @@ class LocalObjectDetector(ObjectDetector):
device_config = {"device": tf_device}
edge_tpu_delegate = None
try:
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:
if tf_device != 'cpu':
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")
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:
print("No EdgeTPU detected. Falling back to CPU.")
@@ -135,9 +132,9 @@ def run_detector(detection_queue, out_events: Dict[str, mp.Event], avg_speed, st
avg_speed.value = (avg_speed.value*9 + duration)/10
class EdgeTPUProcess():
def __init__(self, out_events, tf_device=None):
def __init__(self, detection_queue, out_events, tf_device=None):
self.out_events = out_events
self.detection_queue = mp.Queue()
self.detection_queue = detection_queue
self.avg_inference_speed = mp.Value('d', 0.01)
self.detection_start = mp.Value('d', 0.0)
self.detect_process = None
@@ -192,3 +189,7 @@ class RemoteObjectDetector():
))
self.fps.update()
return detections
def cleanup(self):
self.shm.unlink()
self.out_shm.unlink()

View File

@@ -117,10 +117,9 @@ def start_or_restart_ffmpeg(ffmpeg_cmd, frame_size, ffmpeg_process=None):
def capture_frames(ffmpeg_process, camera_name, frame_shape, frame_manager: FrameManager,
frame_queue, take_frame: int, fps:EventsPerSecond, skipped_fps: EventsPerSecond,
stop_event: mp.Event, detection_frame: mp.Value, current_frame: mp.Value):
stop_event: mp.Event, current_frame: mp.Value):
frame_num = 0
last_frame = 0
frame_size = frame_shape[0] * frame_shape[1] * frame_shape[2]
skipped_fps.start()
while True:
@@ -147,8 +146,8 @@ def capture_frames(ffmpeg_process, camera_name, frame_shape, frame_manager: Fram
skipped_fps.update()
continue
# if the detection process is more than 1 second behind, skip this frame
if detection_frame.value > 0.0 and (last_frame - detection_frame.value) > 1:
# if the queue is full, skip this frame
if frame_queue.full():
skipped_fps.update()
continue
@@ -159,10 +158,9 @@ def capture_frames(ffmpeg_process, camera_name, frame_shape, frame_manager: Fram
# add to the queue
frame_queue.put(current_frame.value)
last_frame = current_frame.value
class CameraCapture(threading.Thread):
def __init__(self, name, ffmpeg_process, frame_shape, frame_queue, take_frame, fps, detection_frame, stop_event):
def __init__(self, name, ffmpeg_process, frame_shape, frame_queue, take_frame, fps, stop_event):
threading.Thread.__init__(self)
self.name = name
self.frame_shape = frame_shape
@@ -175,13 +173,12 @@ class CameraCapture(threading.Thread):
self.ffmpeg_process = ffmpeg_process
self.current_frame = mp.Value('d', 0.0)
self.last_frame = 0
self.detection_frame = detection_frame
self.stop_event = stop_event
def run(self):
self.skipped_fps.start()
capture_frames(self.ffmpeg_process, self.name, self.frame_shape, self.frame_manager, self.frame_queue, self.take_frame,
self.fps, self.skipped_fps, self.stop_event, self.detection_frame, self.current_frame)
self.fps, self.skipped_fps, self.stop_event, self.current_frame)
def track_camera(name, config, frame_queue, frame_shape, detection_queue, result_connection, detected_objects_queue, fps, detection_fps, read_start, detection_frame, stop_event):
print(f"Starting process for {name}: {os.getpid()}")