diff --git a/frigate/detectors/plugins/zmq_ipc.py b/frigate/detectors/plugins/zmq_ipc.py index b65cb801e8..74ce47fbfa 100644 --- a/frigate/detectors/plugins/zmq_ipc.py +++ b/frigate/detectors/plugins/zmq_ipc.py @@ -53,6 +53,12 @@ class ZmqIpcDetector(DetectionApi): "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: @@ -77,6 +83,8 @@ class ZmqIpcDetector(DetectionApi): 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() @@ -148,6 +156,7 @@ class ZmqIpcDetector(DetectionApi): 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 @@ -266,12 +275,26 @@ class ZmqIpcDetector(DetectionApi): 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: diff --git a/frigate/object_detection/base.py b/frigate/object_detection/base.py index 0ab975df6c..49cfc9b3d0 100644 --- a/frigate/object_detection/base.py +++ b/frigate/object_detection/base.py @@ -66,6 +66,10 @@ class BaseLocalDetector(ObjectDetector): self.detect_api.set_stop_event(stop_event) def _transform_input(self, tensor_input: np.ndarray) -> np.ndarray: + # If the detector supports input transform offloading, return the frame unmodified + if getattr(self.detect_api, "input_transform_offload", False) is True: + return tensor_input + if self.input_transform: tensor_input = np.transpose(tensor_input, self.input_transform) diff --git a/frigate/test/test_zmq_input_transform_offload.py b/frigate/test/test_zmq_input_transform_offload.py new file mode 100644 index 0000000000..dad1121621 --- /dev/null +++ b/frigate/test/test_zmq_input_transform_offload.py @@ -0,0 +1,96 @@ +import json +import unittest +from unittest.mock import MagicMock, patch + +import numpy as np +import zmq +from pydantic import TypeAdapter + +import frigate.object_detection.base +from frigate.config import DetectorConfig + +DETECTIONS = [np.zeros((20, 6), np.float32).tobytes()] + + +def handshake(**extra): + return [ + json.dumps({"model_available": True, "model_loaded": True, **extra}).encode() + ] + + +class TestZmqInputTransformOffload(unittest.TestCase): + """The handshake decides whether Frigate applies the input transform. + + The model is configured nchw/float, so the two paths are distinguishable + from the header alone: transformed frames go over as float32 NCHW, + untransformed ones as the uint8 NHWC they were captured in. + """ + + def _headers(self, replies, frames=1): + """Drive a detector through `replies`, return the inference headers.""" + socket = MagicMock() + socket.recv_multipart.side_effect = replies + + with patch("zmq.Context") as context: + context.return_value.socket.return_value = socket + config = TypeAdapter(DetectorConfig).validate_python( + { + "type": "zmq", + "model": { + "path": "/test/model.onnx", + "input_tensor": "nchw", + "input_dtype": "float", + }, + } + ) + detector = frigate.object_detection.base.LocalObjectDetector( + detector_config=config + ) + frame = np.zeros( + (1, config.model.height, config.model.width, 3), dtype=np.uint8 + ) + for _ in range(frames): + detector.detect_raw(frame) + + # send_multipart also carries the handshakes, which are a single frame + return [ + json.loads(call[0][0][0].decode()) + for call in socket.send_multipart.call_args_list + if len(call[0][0]) == 2 + ] + + def test_handshake_decides_the_transform(self): + for reported, dtype, axis, declared in ( + ({"input_transform_offload": True}, "uint8", 3, True), + ({"input_transform_offload": False}, "float32", 1, False), + ({}, "float32", 1, False), + ): + with self.subTest(handshake=reported): + (header,) = self._headers([handshake(**reported), DETECTIONS]) + self.assertEqual(header["dtype"], dtype) + self.assertEqual(header["shape"][axis], 3) + self.assertEqual("input_tensor" in header, declared) + if declared: + self.assertEqual(header["input_tensor"], "nchw") + self.assertEqual(header["input_dtype"], "float") + + def test_transform_follows_a_detector_upgraded_mid_run(self): + """The detector redoes its handshake on a reconnect, so the transform + choice follows it rather than being fixed when Frigate starts.""" + before, after = self._headers( + [ + handshake(), + zmq.Again(), # times out, prompting a reconnect + handshake(input_transform_offload=True), + DETECTIONS, + ], + frames=2, + ) + self.assertEqual(before["dtype"], "float32") + self.assertNotIn("input_tensor", before) + self.assertEqual(after["dtype"], "uint8") + self.assertEqual(after["input_tensor"], "nchw") + + +if __name__ == "__main__": + unittest.main()