|
|
@@ -2,12 +2,14 @@ import os
|
|
|
import datetime
|
|
|
import hashlib
|
|
|
import multiprocessing as mp
|
|
|
+import queue
|
|
|
+from multiprocessing.connection import Connection
|
|
|
from abc import ABC, abstractmethod
|
|
|
+from typing import Dict
|
|
|
import numpy as np
|
|
|
-import pyarrow.plasma as plasma
|
|
|
import tflite_runtime.interpreter as tflite
|
|
|
from tflite_runtime.interpreter import load_delegate
|
|
|
-from frigate.util import EventsPerSecond, listen
|
|
|
+from frigate.util import EventsPerSecond, listen, SharedMemoryFrameManager
|
|
|
|
|
|
def load_labels(path, encoding='utf-8'):
|
|
|
"""Loads labels from file (with or without index numbers).
|
|
|
@@ -100,73 +102,77 @@ class LocalObjectDetector(ObjectDetector):
|
|
|
|
|
|
return detections
|
|
|
|
|
|
-def run_detector(detection_queue, avg_speed, start, tf_device):
|
|
|
+def run_detector(detection_queue, result_connections: Dict[str, Connection], avg_speed, start, tf_device):
|
|
|
print(f"Starting detection process: {os.getpid()}")
|
|
|
listen()
|
|
|
- plasma_client = plasma.connect("/tmp/plasma")
|
|
|
+ frame_manager = SharedMemoryFrameManager()
|
|
|
object_detector = LocalObjectDetector(tf_device=tf_device)
|
|
|
|
|
|
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())
|
|
|
- object_id_out = plasma.ObjectID(hashlib.sha1(str.encode(f"out-{object_id_str}")).digest())
|
|
|
- input_frame = plasma_client.get(object_id, timeout_ms=0)
|
|
|
+ connection_id = detection_queue.get()
|
|
|
+ input_frame = frame_manager.get(connection_id, (1,300,300,3))
|
|
|
|
|
|
- if input_frame is plasma.ObjectNotAvailable:
|
|
|
+ if input_frame is None:
|
|
|
continue
|
|
|
|
|
|
# detect and put the output in the plasma store
|
|
|
start.value = datetime.datetime.now().timestamp()
|
|
|
- plasma_client.put(object_detector.detect_raw(input_frame), object_id_out)
|
|
|
+ # TODO: what is the overhead for pickling this result vs writing back to shared memory?
|
|
|
+ # I could try using an Event() and waiting in the other process before looking in memory...
|
|
|
+ detections = object_detector.detect_raw(input_frame)
|
|
|
+ result_connections[connection_id].send(detections)
|
|
|
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, tf_device=None):
|
|
|
+ def __init__(self, result_connections, tf_device=None):
|
|
|
+ self.result_connections = result_connections
|
|
|
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 stop(self):
|
|
|
+ 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()
|
|
|
|
|
|
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.tf_device))
|
|
|
+ self.stop()
|
|
|
+ self.detect_process = mp.Process(target=run_detector, args=(self.detection_queue, self.result_connections, self.avg_inference_speed, self.detection_start, self.tf_device))
|
|
|
self.detect_process.daemon = True
|
|
|
self.detect_process.start()
|
|
|
|
|
|
class RemoteObjectDetector():
|
|
|
- def __init__(self, name, labels, detection_queue):
|
|
|
+ def __init__(self, name, labels, detection_queue, result_connection: Connection):
|
|
|
self.labels = load_labels(labels)
|
|
|
self.name = name
|
|
|
self.fps = EventsPerSecond()
|
|
|
- self.plasma_client = plasma.connect("/tmp/plasma")
|
|
|
self.detection_queue = detection_queue
|
|
|
+ self.result_connection = result_connection
|
|
|
+ self.shm = mp.shared_memory.SharedMemory(name=self.name, create=True, size=300*300*3)
|
|
|
+ self.np_shm = np.ndarray((1,300,300,3), dtype=np.uint8, buffer=self.shm.buf)
|
|
|
|
|
|
def detect(self, tensor_input, threshold=.4):
|
|
|
detections = []
|
|
|
|
|
|
- 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, timeout_ms=10000)
|
|
|
-
|
|
|
- if raw_detections is plasma.ObjectNotAvailable:
|
|
|
- self.plasma_client.delete([object_id_frame])
|
|
|
+ # copy input to shared memory
|
|
|
+ # TODO: what if I just write it there in the first place?
|
|
|
+ self.np_shm[:] = tensor_input[:]
|
|
|
+ self.detection_queue.put(self.name)
|
|
|
+ if self.result_connection.poll(10):
|
|
|
+ raw_detections = self.result_connection.recv()
|
|
|
+ else:
|
|
|
return detections
|
|
|
|
|
|
for d in raw_detections:
|
|
|
@@ -177,6 +183,5 @@ class RemoteObjectDetector():
|
|
|
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