Add accelerated input transform support to zmq detector protocol (#24731)

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

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:

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)
 

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()