mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-11-13 00:53:02 +03:00
* Let the zmq detector skip the input transform when the detector applies it _transform_input converts the frame to the model's layout and dtype before the detector sees it. For a remote detector that can apply the transform itself, often on a GPU, that work is done twice over and the tensor is sent in its most expensive form. The detector now reports input_transform_offload in the model handshake. When it does, Frigate sends the frame as captured (uint8 NHWC) and adds the model's expected input_tensor and input_dtype to the request header, so the detector can apply them. The model config is unchanged either way: the expected format is relayed, not overridden. The choice is read per frame rather than cached, because the detector redoes its handshake on a timeout or socket error, so restarting or upgrading it while Frigate runs already works and keeps working here. Negotiated rather than configured, for two reasons. The detector is what knows whether it can apply the transform, and a user who merely believed it could would send unconverted frames to a model that cannot take them. It also needs no config surface, which the device string format has no room for. On a 640x640 region, measured on a Ryzen 9 5900XT, the host-side cost of handing the detector a frame drops from about 4 ms to under 0.1 ms per detection, and the payload from 4.9 MB to 1.2 MB. Nearly all of it is tobytes(): np.transpose returns a view and astype keeps its strides by default, so the transpose is deferred to serialization, where it becomes a strided gather over the whole float32 tensor. The YUV to RGB conversion is unaffected and stays on the host. * Test that the handshake decides whether the input transform is applied Parameterized over the three handshake values against the header the detector ends up sending: reported, the frame goes over as captured uint8 NHWC with the model's input_tensor and input_dtype declared; not reported or absent, the existing transform runs and those fields are omitted. A second case drives a real reconnect -- a timed out request, then a handshake from an upgraded detector -- and checks the next frame follows the new answer rather than the one from startup. --------- Co-authored-by: Eric <eric@localhost>
369 lines
14 KiB
Python
369 lines
14 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from typing import Any, ClassVar, Literal
|
|
|
|
import numpy as np
|
|
import zmq
|
|
from pydantic import ConfigDict, Field
|
|
|
|
from frigate.detectors.detection_api import DetectionApi
|
|
from frigate.detectors.detector_config import BaseDetectorConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DETECTOR_KEY = "zmq"
|
|
|
|
|
|
class ZmqDetectorConfig(BaseDetectorConfig):
|
|
"""ZMQ IPC detector that offloads inference to an external process via a ZeroMQ IPC endpoint."""
|
|
|
|
model_config = ConfigDict(
|
|
title="ZMQ IPC",
|
|
)
|
|
device_spec_field: ClassVar[str] = "endpoint"
|
|
|
|
type: Literal[DETECTOR_KEY]
|
|
endpoint: str = Field(
|
|
default="ipc:///tmp/cache/zmq_detector",
|
|
title="ZMQ IPC endpoint",
|
|
description="The ZMQ endpoint to connect to.",
|
|
)
|
|
request_timeout_ms: int = Field(
|
|
default=200,
|
|
title="ZMQ request timeout in milliseconds",
|
|
description="Timeout for ZMQ requests in milliseconds.",
|
|
)
|
|
linger_ms: int = Field(
|
|
default=0,
|
|
title="ZMQ socket linger in milliseconds",
|
|
description="Socket linger period in milliseconds.",
|
|
)
|
|
|
|
|
|
class ZmqIpcDetector(DetectionApi):
|
|
"""
|
|
ZMQ-based detector plugin using a REQ/REP socket over an IPC endpoint.
|
|
|
|
Protocol:
|
|
- Request is sent as a multipart message:
|
|
[ header_json_bytes, tensor_bytes ]
|
|
where header is a JSON object containing:
|
|
{
|
|
"shape": List[int],
|
|
"dtype": str, # numpy dtype string, e.g. "uint8", "float32"
|
|
}
|
|
and, when the detector reported input_transform_offload, also:
|
|
{
|
|
"input_tensor": str, # InputTensorEnum value, e.g. "nhwc", "nchw"
|
|
"input_dtype": str, # InputDTypeEnum value: "float", "float_denorm", "int"
|
|
}
|
|
which describe the transform the detector is expected to apply itself.
|
|
tensor_bytes are the raw bytes of the numpy array in C-order.
|
|
|
|
- Response is expected to be either:
|
|
a) Multipart [ header_json_bytes, tensor_bytes ] with header specifying
|
|
shape [20,6] and dtype "float32"; or
|
|
b) Single frame tensor_bytes of length 20*6*4 bytes (float32).
|
|
|
|
On any error or timeout, this detector returns a zero array of shape (20, 6).
|
|
|
|
Model Management:
|
|
- On initialization, sends model request to check if model is available
|
|
- If model not available, sends model data via ZMQ
|
|
- Only starts inference after model is ready
|
|
"""
|
|
|
|
type_key = DETECTOR_KEY
|
|
|
|
def __init__(self, detector_config: ZmqDetectorConfig):
|
|
super().__init__(detector_config)
|
|
|
|
self._context = zmq.Context()
|
|
self._endpoint = detector_config.endpoint
|
|
self._request_timeout_ms = detector_config.request_timeout_ms
|
|
self._linger_ms = detector_config.linger_ms
|
|
# Reported by the detector during the model handshake.
|
|
self.input_transform_offload = False
|
|
self._socket = None
|
|
self._create_socket()
|
|
|
|
# Model management
|
|
self._model_ready = False
|
|
self._model_name = self._get_model_name()
|
|
|
|
# Initialize model if needed
|
|
self._initialize_model()
|
|
|
|
# Preallocate zero result for error paths
|
|
self._zero_result = np.zeros((20, 6), np.float32)
|
|
|
|
def _create_socket(self) -> None:
|
|
if self._socket is not None:
|
|
try:
|
|
self._socket.close(linger=self._linger_ms)
|
|
except Exception:
|
|
pass
|
|
self._socket = self._context.socket(zmq.REQ)
|
|
# Apply timeouts and linger so calls don't block indefinitely
|
|
self._socket.setsockopt(zmq.RCVTIMEO, self._request_timeout_ms)
|
|
self._socket.setsockopt(zmq.SNDTIMEO, self._request_timeout_ms)
|
|
self._socket.setsockopt(zmq.LINGER, self._linger_ms)
|
|
|
|
logger.debug(f"ZMQ detector connecting to {self._endpoint}")
|
|
self._socket.connect(self._endpoint)
|
|
|
|
def _get_model_name(self) -> str:
|
|
"""Get the model filename from the detector config."""
|
|
model_path = self.detector_config.model.path
|
|
return os.path.basename(model_path)
|
|
|
|
def _initialize_model(self) -> None:
|
|
"""Initialize the model by checking availability and transferring if needed."""
|
|
try:
|
|
logger.info(f"Initializing model: {self._model_name}")
|
|
|
|
# Check if model is available and transfer if needed
|
|
if self._check_and_transfer_model():
|
|
logger.info(f"Model {self._model_name} is ready")
|
|
self._model_ready = True
|
|
else:
|
|
logger.error(f"Failed to initialize model {self._model_name}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize model: {e}")
|
|
|
|
def _check_and_transfer_model(self) -> bool:
|
|
"""Check if model is available and transfer if needed in one atomic operation."""
|
|
try:
|
|
# Send model availability request
|
|
header = {"model_request": True, "model_name": self._model_name}
|
|
header_bytes = json.dumps(header).encode("utf-8")
|
|
|
|
self._socket.send_multipart([header_bytes])
|
|
|
|
# Temporarily increase timeout for model operations
|
|
original_timeout = self._socket.getsockopt(zmq.RCVTIMEO)
|
|
self._socket.setsockopt(zmq.RCVTIMEO, 30000)
|
|
|
|
try:
|
|
response_frames = self._socket.recv_multipart()
|
|
finally:
|
|
self._socket.setsockopt(zmq.RCVTIMEO, original_timeout)
|
|
|
|
if len(response_frames) == 1:
|
|
try:
|
|
response = json.loads(response_frames[0].decode("utf-8"))
|
|
model_available = response.get("model_available", False)
|
|
model_loaded = response.get("model_loaded", False)
|
|
self._read_offload_support(response)
|
|
|
|
if model_available and model_loaded:
|
|
return True
|
|
elif model_available and not model_loaded:
|
|
logger.error("Model exists but failed to load")
|
|
return False
|
|
else:
|
|
return self._send_model_data()
|
|
|
|
except json.JSONDecodeError:
|
|
logger.warning(
|
|
"Received non-JSON response for model availability check"
|
|
)
|
|
return False
|
|
else:
|
|
logger.warning(
|
|
"Received unexpected response format for model availability check"
|
|
)
|
|
return False
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to check and transfer model: {e}")
|
|
return False
|
|
|
|
def _check_model_availability(self) -> bool:
|
|
"""Check if the model is available on the detector."""
|
|
try:
|
|
# Send model availability request
|
|
header = {"model_request": True, "model_name": self._model_name}
|
|
header_bytes = json.dumps(header).encode("utf-8")
|
|
|
|
self._socket.send_multipart([header_bytes])
|
|
|
|
# Receive response
|
|
response_frames = self._socket.recv_multipart()
|
|
|
|
# Check if this is a JSON response (model management)
|
|
if len(response_frames) == 1:
|
|
try:
|
|
response = json.loads(response_frames[0].decode("utf-8"))
|
|
model_available = response.get("model_available", False)
|
|
model_loaded = response.get("model_loaded", False)
|
|
logger.debug(
|
|
f"Model availability check: available={model_available}, loaded={model_loaded}"
|
|
)
|
|
return model_available and model_loaded
|
|
except json.JSONDecodeError:
|
|
logger.warning(
|
|
"Received non-JSON response for model availability check"
|
|
)
|
|
return False
|
|
else:
|
|
logger.warning(
|
|
"Received unexpected response format for model availability check"
|
|
)
|
|
return False
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to check model availability: {e}")
|
|
return False
|
|
|
|
def _send_model_data(self) -> bool:
|
|
"""Send model data to the detector."""
|
|
try:
|
|
model_path = self.detector_config.model.path
|
|
|
|
if not os.path.exists(model_path):
|
|
logger.error(f"Model file not found: {model_path}")
|
|
return False
|
|
|
|
logger.info(f"Transferring model to detector: {self._model_name}")
|
|
with open(model_path, "rb") as f:
|
|
model_data = f.read()
|
|
|
|
header = {"model_data": True, "model_name": self._model_name}
|
|
header_bytes = json.dumps(header).encode("utf-8")
|
|
|
|
self._socket.send_multipart([header_bytes, model_data])
|
|
|
|
# Temporarily increase timeout for model loading (can take several seconds)
|
|
original_timeout = self._socket.getsockopt(zmq.RCVTIMEO)
|
|
self._socket.setsockopt(zmq.RCVTIMEO, 30000)
|
|
|
|
try:
|
|
# Receive response
|
|
response_frames = self._socket.recv_multipart()
|
|
finally:
|
|
# Restore original timeout
|
|
self._socket.setsockopt(zmq.RCVTIMEO, original_timeout)
|
|
|
|
# Check if this is a JSON response (model management)
|
|
if len(response_frames) == 1:
|
|
try:
|
|
response = json.loads(response_frames[0].decode("utf-8"))
|
|
model_saved = response.get("model_saved", False)
|
|
model_loaded = response.get("model_loaded", False)
|
|
if model_saved and model_loaded:
|
|
logger.info(
|
|
f"Model {self._model_name} transferred and loaded successfully"
|
|
)
|
|
else:
|
|
logger.error(
|
|
f"Model transfer failed: saved={model_saved}, loaded={model_loaded}"
|
|
)
|
|
return model_saved and model_loaded
|
|
except json.JSONDecodeError:
|
|
logger.warning("Received non-JSON response for model data transfer")
|
|
return False
|
|
else:
|
|
logger.warning(
|
|
"Received unexpected response format for model data transfer"
|
|
)
|
|
return False
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to send model data: {e}")
|
|
return False
|
|
|
|
def _read_offload_support(self, response: dict[str, Any]) -> None:
|
|
"""Record whether the detector applies the input transform itself."""
|
|
offload = response.get("input_transform_offload", False) is True
|
|
if offload != self.input_transform_offload:
|
|
self.input_transform_offload = offload
|
|
logger.info(
|
|
f"Detector input transform offload: {'on' if offload else 'off'}"
|
|
)
|
|
|
|
def _build_header(self, tensor_input: np.ndarray) -> bytes:
|
|
header: dict[str, Any] = {
|
|
"shape": list(tensor_input.shape),
|
|
"dtype": str(tensor_input.dtype.name),
|
|
"model_type": str(self.detector_config.model.model_type.name),
|
|
}
|
|
|
|
if self.input_transform_offload:
|
|
header["input_tensor"] = self.detector_config.model.input_tensor.value
|
|
header["input_dtype"] = self.detector_config.model.input_dtype.value
|
|
|
|
return json.dumps(header).encode("utf-8")
|
|
|
|
def _decode_response(self, frames: list[bytes]) -> np.ndarray:
|
|
try:
|
|
if len(frames) == 1:
|
|
# Single-frame raw float32 (20x6)
|
|
buf = frames[0]
|
|
if len(buf) != 20 * 6 * 4:
|
|
logger.warning(
|
|
f"ZMQ detector received unexpected payload size: {len(buf)}"
|
|
)
|
|
return self._zero_result
|
|
return np.frombuffer(buf, dtype=np.float32).reshape((20, 6))
|
|
|
|
if len(frames) >= 2:
|
|
header = json.loads(frames[0].decode("utf-8"))
|
|
shape = tuple(header.get("shape", []))
|
|
dtype = np.dtype(header.get("dtype", "float32"))
|
|
return np.frombuffer(frames[1], dtype=dtype).reshape(shape)
|
|
|
|
logger.warning("ZMQ detector received empty reply")
|
|
return self._zero_result
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.error(f"ZMQ detector failed to decode response: {exc}")
|
|
return self._zero_result
|
|
|
|
def detect_raw(self, tensor_input: np.ndarray) -> np.ndarray:
|
|
if not self._model_ready:
|
|
logger.warning("Model not ready, returning zero detections")
|
|
return self._zero_result
|
|
|
|
try:
|
|
header_bytes = self._build_header(tensor_input)
|
|
payload_bytes = memoryview(tensor_input.tobytes(order="C"))
|
|
|
|
# Send request
|
|
self._socket.send_multipart([header_bytes, payload_bytes])
|
|
|
|
# Receive reply
|
|
reply_frames = self._socket.recv_multipart()
|
|
detections = self._decode_response(reply_frames)
|
|
|
|
# Ensure output shape and dtype are exactly as expected
|
|
return detections
|
|
except zmq.Again:
|
|
# Timeout
|
|
logger.debug("ZMQ detector request timed out; resetting socket")
|
|
try:
|
|
self._create_socket()
|
|
self._initialize_model()
|
|
except Exception:
|
|
pass
|
|
return self._zero_result
|
|
except zmq.ZMQError as exc:
|
|
logger.error(f"ZMQ detector ZMQError: {exc}; resetting socket")
|
|
try:
|
|
self._create_socket()
|
|
self._initialize_model()
|
|
except Exception:
|
|
pass
|
|
return self._zero_result
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.error(f"ZMQ detector unexpected error: {exc}")
|
|
return self._zero_result
|
|
|
|
def __del__(self) -> None: # pragma: no cover - best-effort cleanup
|
|
try:
|
|
if self._socket is not None:
|
|
self._socket.close(linger=self.detector_config.linger_ms)
|
|
except Exception:
|
|
pass
|