mirror of
https://github.com/alexhopeoconnor/frigate.git
synced 2026-10-22 03:02:09 +10:00
Improve async object detector support (#17712)
* Move object detection to folder * Add input store type * Add hwnc * Add hwcn * Fix test
This commit is contained in:
@@ -0,0 +1,77 @@
|
||||
"""Object detection utilities."""
|
||||
|
||||
import queue
|
||||
import threading
|
||||
|
||||
from numpy import ndarray
|
||||
|
||||
from frigate.detectors.detector_config import InputTensorEnum
|
||||
|
||||
|
||||
class RequestStore:
|
||||
"""
|
||||
A thread-safe hash-based response store that handles creating requests.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.request_counter = 0
|
||||
self.request_counter_lock = threading.Lock()
|
||||
self.input_queue = queue.Queue()
|
||||
|
||||
def __get_request_id(self) -> int:
|
||||
with self.request_counter_lock:
|
||||
request_id = self.request_counter
|
||||
self.request_counter += 1
|
||||
if self.request_counter > 1000000:
|
||||
self.request_counter = 0
|
||||
return request_id
|
||||
|
||||
def put(self, tensor_input: ndarray) -> int:
|
||||
request_id = self.__get_request_id()
|
||||
self.input_queue.get((request_id, tensor_input))
|
||||
return request_id
|
||||
|
||||
def get(self) -> tuple[int, ndarray] | None:
|
||||
try:
|
||||
return self.input_queue.get_nowait()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class ResponseStore:
|
||||
"""
|
||||
A thread-safe hash-based response store that maps request IDs
|
||||
to their results. Threads can wait on the condition variable until
|
||||
their request's result appears.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.responses = {} # Maps request_id -> (original_input, infer_results)
|
||||
self.lock = threading.Lock()
|
||||
self.cond = threading.Condition(self.lock)
|
||||
|
||||
def put(self, request_id: int, response: ndarray):
|
||||
with self.cond:
|
||||
self.responses[request_id] = response
|
||||
self.cond.notify_all()
|
||||
|
||||
def get(self, request_id: int, timeout=None) -> ndarray:
|
||||
with self.cond:
|
||||
if not self.cond.wait_for(
|
||||
lambda: request_id in self.responses, timeout=timeout
|
||||
):
|
||||
raise TimeoutError(f"Timeout waiting for response {request_id}")
|
||||
|
||||
return self.responses.pop(request_id)
|
||||
|
||||
|
||||
def tensor_transform(desired_shape: InputTensorEnum):
|
||||
# Currently this function only supports BHWC permutations
|
||||
if desired_shape == InputTensorEnum.nhwc:
|
||||
return None
|
||||
elif desired_shape == InputTensorEnum.nchw:
|
||||
return (0, 3, 1, 2)
|
||||
elif desired_shape == InputTensorEnum.hwnc:
|
||||
return (1, 2, 0, 3)
|
||||
elif desired_shape == InputTensorEnum.hwcn:
|
||||
return (1, 2, 3, 0)
|
||||
Reference in New Issue
Block a user