chris/frigate
1
0
Fork 0
forked from Github/frigate
Code Pull Requests Packages Activity

upgrade to python3.8 and switch from plasma store to shared_memory

Browse Source
This commit is contained in:
Blake Blackshear
2020-09-21 21:02:00 -05:00
parent b063099b2a
commit ec4d048905
7 changed files with 128 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
frigate/edgetpu.py
View File

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