improve detection processing and restart when stuck

frigate/edgetpu.py

@@ -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
 
+def run_detector(detection_queue, avg_speed, start):
+    print(f"Starting detection process: {os.getpid()}")
+    plasma_client = plasma.connect("/tmp/plasma")
+    object_detector = ObjectDetector()
+
+    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)
+
+        start.value = datetime.datetime.now().timestamp()
+
+        # 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))
+
+        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):
-        # 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
-
-        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)
-
-        self.detect_lock = mp.Lock()
-        self.detect_ready = mp.Event()
-        self.frame_ready = mp.Event()
+        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()
 
-        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")
-
-            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
-
-        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