|
|
@@ -1,8 +1,10 @@
|
|
|
import os
|
|
|
import datetime
|
|
|
+import hashlib
|
|
|
import multiprocessing as mp
|
|
|
import numpy as np
|
|
|
import SharedArray as sa
|
|
|
+import pyarrow.plasma as plasma
|
|
|
import tflite_runtime.interpreter as tflite
|
|
|
from tflite_runtime.interpreter import load_delegate
|
|
|
from frigate.util import EventsPerSecond
|
|
|
@@ -60,77 +62,75 @@ class ObjectDetector():
|
|
|
|
|
|
return detections
|
|
|
|
|
|
-class EdgeTPUProcess():
|
|
|
- def __init__(self):
|
|
|
- # TODO: see if we can use the plasma store with a queue and maintain the same speeds
|
|
|
- try:
|
|
|
- sa.delete("frame")
|
|
|
- except:
|
|
|
- pass
|
|
|
- try:
|
|
|
- sa.delete("detections")
|
|
|
- except:
|
|
|
- pass
|
|
|
+def run_detector(detection_queue, avg_speed, start):
|
|
|
+ print(f"Starting detection process: {os.getpid()}")
|
|
|
+ plasma_client = plasma.connect("/tmp/plasma")
|
|
|
+ object_detector = ObjectDetector()
|
|
|
|
|
|
- self.input_frame = sa.create("frame", shape=(1,300,300,3), dtype=np.uint8)
|
|
|
- self.detections = sa.create("detections", shape=(20,6), dtype=np.float32)
|
|
|
+ while True:
|
|
|
+ object_id_str = detection_queue.get()
|
|
|
+ object_id_hash = hashlib.sha1(str.encode(object_id_str))
|
|
|
+ object_id = plasma.ObjectID(object_id_hash.digest())
|
|
|
+ input_frame = plasma_client.get(object_id, timeout_ms=0)
|
|
|
|
|
|
- self.detect_lock = mp.Lock()
|
|
|
- self.detect_ready = mp.Event()
|
|
|
- self.frame_ready = mp.Event()
|
|
|
- self.avg_inference_speed = mp.Value('d', 0.01)
|
|
|
+ start.value = datetime.datetime.now().timestamp()
|
|
|
|
|
|
- def run_detector(detect_ready, frame_ready, avg_speed):
|
|
|
- print(f"Starting detection process: {os.getpid()}")
|
|
|
- object_detector = ObjectDetector()
|
|
|
- input_frame = sa.attach("frame")
|
|
|
- detections = sa.attach("detections")
|
|
|
+ # detect and put the output in the plasma store
|
|
|
+ object_id_out = hashlib.sha1(str.encode(f"out-{object_id_str}")).digest()
|
|
|
+ plasma_client.put(object_detector.detect_raw(input_frame), plasma.ObjectID(object_id_out))
|
|
|
|
|
|
- while True:
|
|
|
- # wait until a frame is ready
|
|
|
- frame_ready.wait()
|
|
|
- start = datetime.datetime.now().timestamp()
|
|
|
- # signal that the process is busy
|
|
|
- frame_ready.clear()
|
|
|
- detections[:] = object_detector.detect_raw(input_frame)
|
|
|
- # signal that the process is ready to detect
|
|
|
- detect_ready.set()
|
|
|
- duration = datetime.datetime.now().timestamp()-start
|
|
|
- avg_speed.value = (avg_speed.value*9 + duration)/10
|
|
|
+ duration = datetime.datetime.now().timestamp()-start.value
|
|
|
+ start.value = 0.0
|
|
|
+ avg_speed.value = (avg_speed.value*9 + duration)/10
|
|
|
+
|
|
|
+class EdgeTPUProcess():
|
|
|
+ def __init__(self):
|
|
|
+ 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.start_or_restart()
|
|
|
|
|
|
- self.detect_process = mp.Process(target=run_detector, args=(self.detect_ready, self.frame_ready, self.avg_inference_speed))
|
|
|
+ def start_or_restart(self):
|
|
|
+ self.detection_start.value = 0.0
|
|
|
+ if (not self.detect_process is None) and self.detect_process.is_alive():
|
|
|
+ self.detect_process.terminate()
|
|
|
+ print("Waiting for detection process to exit gracefully...")
|
|
|
+ self.detect_process.join(timeout=30)
|
|
|
+ if self.detect_process.exitcode is None:
|
|
|
+ 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.daemon = True
|
|
|
self.detect_process.start()
|
|
|
|
|
|
class RemoteObjectDetector():
|
|
|
- def __init__(self, labels, detect_lock, detect_ready, frame_ready):
|
|
|
+ def __init__(self, name, labels, detection_queue):
|
|
|
self.labels = load_labels(labels)
|
|
|
-
|
|
|
- self.input_frame = sa.attach("frame")
|
|
|
- self.detections = sa.attach("detections")
|
|
|
-
|
|
|
+ self.name = name
|
|
|
self.fps = EventsPerSecond()
|
|
|
-
|
|
|
- self.detect_lock = detect_lock
|
|
|
- self.detect_ready = detect_ready
|
|
|
- self.frame_ready = frame_ready
|
|
|
+ self.plasma_client = plasma.connect("/tmp/plasma")
|
|
|
+ self.detection_queue = detection_queue
|
|
|
|
|
|
def detect(self, tensor_input, threshold=.4):
|
|
|
detections = []
|
|
|
- with self.detect_lock:
|
|
|
- self.input_frame[:] = tensor_input
|
|
|
- # unset detections and signal that a frame is ready
|
|
|
- self.detect_ready.clear()
|
|
|
- self.frame_ready.set()
|
|
|
- # wait until the detection process is finished,
|
|
|
- self.detect_ready.wait()
|
|
|
- for d in self.detections:
|
|
|
- if d[1] < threshold:
|
|
|
- break
|
|
|
- detections.append((
|
|
|
- self.labels[int(d[0])],
|
|
|
- float(d[1]),
|
|
|
- (d[2], d[3], d[4], d[5])
|
|
|
- ))
|
|
|
+
|
|
|
+ now = f"{self.name}-{str(datetime.datetime.now().timestamp())}"
|
|
|
+ object_id_frame = plasma.ObjectID(hashlib.sha1(str.encode(now)).digest())
|
|
|
+ object_id_detections = plasma.ObjectID(hashlib.sha1(str.encode(f"out-{now}")).digest())
|
|
|
+ self.plasma_client.put(tensor_input, object_id_frame)
|
|
|
+ self.detection_queue.put(now)
|
|
|
+ raw_detections = self.plasma_client.get(object_id_detections)
|
|
|
+
|
|
|
+ for d in raw_detections:
|
|
|
+ if d[1] < threshold:
|
|
|
+ break
|
|
|
+ detections.append((
|
|
|
+ self.labels[int(d[0])],
|
|
|
+ float(d[1]),
|
|
|
+ (d[2], d[3], d[4], d[5])
|
|
|
+ ))
|
|
|
+ self.plasma_client.delete([object_id_frame, object_id_detections])
|
|
|
self.fps.update()
|
|
|
return detections
|
Blake Blackshear