Adds FrameStream.overlay_with for fork-join pipelines

examples/blob_tracker_overlay.py

@@ -0,0 +1,75 @@
+"""
+Fork-join blob tracker: render events to frames, run a tracker on the same events,
+and overlay the tracker output on the rendered frames.
+
+                     |-> render ------------------> frames -|
+    events -> regularize                                    |-> overlay_with(draw) -> video
+                     |-> run_tracker -> tracker positions --|
+
+Both branches are built from the same `regularize(frequency_hz=..., start=...)`
+parameters, so packet i in each branch covers the same time window and the merge is a
+positional zip (see `FrameStream.overlay_with`).
+
+`run_tracker` is the boundary to an arbitrary backend. Here it is a trivial Python
+placeholder; swap it for Greg's C++ tracker (the packet is a zero-copy structured-array
+view -- hand its buffer to C++ via pybind11/DLPack and return the detections).
+"""
+
+import numpy
+
+import faery
+
+INPUT = faery.dirname.parent / "tests" / "data" / "dvs.es"
+OUTPUT = faery.dirname.parent / "tests" / "data_generated" / "dvs_tracked.mp4"
+FREQUENCY_HZ = 60.0
+START = 0 * faery.us
+
+
+# --- tracker branch --------------------------------------------------------------
+# PLACEHOLDER. Replace the body with your detector (e.g. call into C++). It must return,
+# per packet, an iterable of (x, y, radius) detections. Returning [] means "no object".
+def run_tracker(packet: numpy.ndarray) -> list[tuple[int, int, int]]:
+    if len(packet) == 0:
+        return []
+    x = int(round(packet["x"].mean()))
+    y = int(round(packet["y"].mean()))
+    spread = int(round(numpy.hypot(packet["x"].std(), packet["y"].std()))) + 5
+    return [(x, y, spread)]
+
+
+# --- draw / merge branch ---------------------------------------------------------
+# Plotting only: rasterize hollow circles for each detection onto the RGBA frame.
+def draw_circles(
+    pixels: numpy.ndarray,
+    detections: list[tuple[int, int, int]],
+    color=(255, 255, 0, 255),
+    thickness: float = 1.5,
+) -> None:
+    height, width = pixels.shape[:2]
+    color_array = numpy.array(color, dtype=numpy.uint8)
+    yy, xx = numpy.mgrid[0:height, 0:width]
+    for cx, cy, radius in detections:
+        distance = numpy.hypot(xx - cx, yy - cy)
+        ring = numpy.abs(distance - radius) <= thickness
+        pixels[ring] = color_array
+
+
+def main() -> None:
+    source = faery.events_stream_from_file(INPUT)
+
+    frames = source.regularize(frequency_hz=FREQUENCY_HZ, start=START).render(
+        decay="exponential",
+        tau="00:00:00.200000",
+        colormap=faery.colormaps.managua.flipped(),
+    )
+    tracks = (
+        run_tracker(packet)
+        for packet in source.regularize(frequency_hz=FREQUENCY_HZ, start=START)
+    )
+
+    frames.overlay_with(tracks, draw=draw_circles).to_file(OUTPUT)
+    print(f"wrote {OUTPUT}")
+
+
+if __name__ == "__main__":
+    main()

python/faery/frame_filter.py

@@ -302,3 +302,70 @@ def __iter__(self) -> collections.abc.Iterator[frame_stream.Frame]:
                 pixels=self.function(frame.pixels),
                 t=frame.t,
             )
+
+
+@typed_filter({"", "Finite", "Regular", "FiniteRegular"})
+class Overlay(frame_stream.FiniteRegularFrameFilter):
+    """
+    Draws sidecar data onto each frame by zipping the frame stream with a second,
+    index-aligned iterable.
+
+    This is the "merge" step of a fork-join pipeline: a source event stream is used
+    twice, once to `render` frames and once to compute per-packet data (e.g. tracker
+    positions), and this filter recombines the two branches.
+
+        source = events.regularize(frequency_hz=hz)
+        frames = source.render(...)
+        tracks = (run_tracker(packet) for packet in events.regularize(frequency_hz=hz))
+        video = frames.overlay(tracks, draw=draw_circles)
+
+    Alignment is positional: sidecar item `i` is paired with frame `i`. Both branches
+    are assumed to share the same packet cadence -- the simplest guarantee is to build
+    both from the same `regularize(frequency_hz=..., start=...)` parameters. Pin `start`
+    if the two branches do not share a parent object.
+
+    `run_tracker` (the sidecar producer) and `draw` are arbitrary callables. `run_tracker`
+    is the natural boundary to a foreign backend: it receives a numpy structured-array
+    view of the packet (zero-copy, DLPack-compatible) and may forward it to C++/Rust,
+    returning any Python object. `draw` receives the RGBA frame and the sidecar item.
+
+    Args:
+        parent: Input frame stream.
+        sidecar: Iterable yielding one item per frame, aligned by index. Must be at
+            least as long as the frame stream.
+        draw: Callable `(pixels, item)` invoked per frame. `pixels` is the RGBA frame
+            (shape `(height, width, 4)`, dtype uint8). Mutate it in place, or return a
+            replacement array; returning `None` keeps the in-place result.
+
+    Raises:
+        ValueError: If the sidecar is exhausted before the frame stream (cadence mismatch).
+    """
+
+    def __init__(
+        self,
+        parent: frame_stream.FiniteRegularFrameStream,
+        sidecar: collections.abc.Iterable[typing.Any],
+        draw: collections.abc.Callable[
+            [numpy.typing.NDArray[numpy.uint8], typing.Any],
+            typing.Optional[numpy.typing.NDArray[numpy.uint8]],
+        ],
+    ):
+        self.init(parent=parent)
+        self.sidecar = sidecar
+        self.draw = draw
+
+    def __iter__(self) -> collections.abc.Iterator[frame_stream.Frame]:
+        sidecar_iterator = iter(self.sidecar)
+        for index, frame in enumerate(self.parent):
+            try:
+                item = next(sidecar_iterator)
+            except StopIteration:
+                raise ValueError(
+                    f"the overlay sidecar was exhausted after {index} item(s) but the "
+                    "frame stream produced more frames; the two branches are not aligned "
+                    "(ensure both use the same regularize frequency_hz and start)"
+                )
+            replacement = self.draw(frame.pixels, item)
+            if replacement is not None:
+                frame.pixels = replacement
+            yield frame

python/faery/frame_stream.py

@@ -180,6 +180,15 @@ def map(
         ],
     ) -> "FrameStream": ...
 
+    def overlay_with(
+        self,
+        sidecar: collections.abc.Iterable[typing.Any],
+        draw: collections.abc.Callable[
+            [numpy.typing.NDArray[numpy.uint8], typing.Any],
+            typing.Optional[numpy.typing.NDArray[numpy.uint8]],
+        ],
+    ) -> "FrameStream": ...
+
 
 class FiniteFrameStream(
     stream.FiniteStream[Frame],
@@ -224,6 +233,15 @@ def map(
         ],
     ) -> "FiniteFrameStream": ...
 
+    def overlay_with(
+        self,
+        sidecar: collections.abc.Iterable[typing.Any],
+        draw: collections.abc.Callable[
+            [numpy.typing.NDArray[numpy.uint8], typing.Any],
+            typing.Optional[numpy.typing.NDArray[numpy.uint8]],
+        ],
+    ) -> "FiniteFrameStream": ...
+
 
 class RegularFrameStream(
     stream.RegularStream[Frame],
@@ -269,6 +287,15 @@ def map(
         ],
     ) -> "RegularFrameStream": ...
 
+    def overlay_with(
+        self,
+        sidecar: collections.abc.Iterable[typing.Any],
+        draw: collections.abc.Callable[
+            [numpy.typing.NDArray[numpy.uint8], typing.Any],
+            typing.Optional[numpy.typing.NDArray[numpy.uint8]],
+        ],
+    ) -> "RegularFrameStream": ...
+
 
 class FiniteRegularFrameStream(
     stream.FiniteRegularStream[Frame],
@@ -314,6 +341,15 @@ def map(
         ],
     ) -> "FiniteRegularFrameStream": ...
 
+    def overlay_with(
+        self,
+        sidecar: collections.abc.Iterable[typing.Any],
+        draw: collections.abc.Callable[
+            [numpy.typing.NDArray[numpy.uint8], typing.Any],
+            typing.Optional[numpy.typing.NDArray[numpy.uint8]],
+        ],
+    ) -> "FiniteRegularFrameStream": ...
+
 
 def bind(prefix: typing.Literal["", "Finite", "Regular", "FiniteRegular"]):
 
@@ -426,13 +462,32 @@ def map(
             function=function,
         )
 
+    def overlay_with(
+        self,
+        sidecar: collections.abc.Iterable[typing.Any],
+        draw: collections.abc.Callable[
+            [numpy.typing.NDArray[numpy.uint8], typing.Any],
+            typing.Optional[numpy.typing.NDArray[numpy.uint8]],
+        ],
+    ):
+        from .frame_filter import FILTERS
+
+        return FILTERS[f"{prefix}Overlay"](
+            parent=self,
+            sidecar=sidecar,
+            draw=draw,
+        )
+
     scale.filter_return_annotation = f"{prefix}FrameStream"
     annotate.filter_return_annotation = f"{prefix}FrameStream"
+    map.filter_return_annotation = f"{prefix}FrameStream"
+    overlay_with.filter_return_annotation = f"{prefix}FrameStream"
 
     globals()[f"{prefix}FrameStream"].scale = scale
     globals()[f"{prefix}FrameStream"].annotate = annotate
     globals()[f"{prefix}FrameStream"].add_overlay = add_overlay
     globals()[f"{prefix}FrameStream"].map = map
+    globals()[f"{prefix}FrameStream"].overlay_with = overlay_with
 
 
 for prefix in ("", "Finite", "Regular", "FiniteRegular"):

tests/test_overlay.py

@@ -0,0 +1,109 @@
+"""
+Tests for `FrameStream.overlay_with`, the "merge" step of a fork-join pipeline:
+one event stream is used twice -- once to render frames, once to compute per-packet
+sidecar data (e.g. tracker positions) -- and the two branches are recombined by
+drawing the sidecar onto the frames.
+"""
+
+import numpy
+import pytest
+
+import faery
+
+EVENTS = numpy.array(
+    [
+        (100, 5, 5, True),  # window 0 -> [0, 1000)
+        (1100, 10, 8, True),  # window 1 -> [1000, 2000)
+        # window 2 -> [2000, 3000) is intentionally empty
+        (3100, 20, 20, False),  # window 3 -> [3000, 4000)
+    ],
+    dtype=faery.EVENTS_DTYPE,
+)
+DIMENSIONS = (32, 24)  # width, height
+FREQUENCY_HZ = 1000.0  # period = 1000 us, so one window per event gap
+
+RED = numpy.array([255, 0, 0, 255], dtype=numpy.uint8)
+
+
+def centroid(packet: numpy.ndarray):
+    """Toy 'tracker': returns the (x, y) centroid of a packet, or None if empty.
+
+    This stands in for a foreign backend. `packet` is a zero-copy structured-array
+    view, exactly what would be handed to a C++/Rust tracker through pybind/DLPack.
+    """
+    if len(packet) == 0:
+        return None
+    return (int(round(packet["x"].mean())), int(round(packet["y"].mean())))
+
+
+def draw_marker(pixels: numpy.ndarray, position) -> None:
+    """Draws the tracker position onto the frame in place (returns None)."""
+    if position is None:
+        return
+    x, y = position
+    pixels[y, x] = RED
+
+
+def branches():
+    source = faery.events_stream_from_array(EVENTS, dimensions=DIMENSIONS)
+    # Two branches off the same source, pinned to the same regularize parameters so
+    # that packet i in each branch covers the identical time window (index alignment).
+    frames = source.regularize(frequency_hz=FREQUENCY_HZ, start=0 * faery.us).render(
+        decay="exponential",
+        tau="00:00:00.001000",
+        colormap=faery.colormaps.managua,
+    )
+    tracks = (
+        centroid(packet)
+        for packet in source.regularize(frequency_hz=FREQUENCY_HZ, start=0 * faery.us)
+    )
+    return frames, tracks
+
+
+def test_overlay_aligns_and_draws():
+    frames, tracks = branches()
+    video = frames.overlay_with(tracks, draw=draw_marker)
+
+    rendered = list(video)
+
+    # Independently recompute the expected per-packet centroids from the same source,
+    # so a match confirms frame i was paired with the sidecar derived from packet i.
+    expected = [
+        centroid(packet)
+        for packet in faery.events_stream_from_array(
+            EVENTS, dimensions=DIMENSIONS
+        ).regularize(frequency_hz=FREQUENCY_HZ, start=0 * faery.us)
+    ]
+
+    assert len(rendered) == len(expected)
+    assert expected == [(5, 5), (10, 8), None, (20, 20)]
+
+    for frame, position in zip(rendered, expected):
+        if position is None:
+            # empty packet -> nothing drawn -> no pure-red marker anywhere
+            assert not numpy.any(numpy.all(frame.pixels == RED, axis=-1))
+        else:
+            x, y = position
+            assert numpy.array_equal(frame.pixels[y, x], RED)
+
+
+def test_overlay_replacement_return_value():
+    """A draw callback that returns an array replaces the frame's pixels."""
+    frames, tracks = branches()
+
+    replacement = numpy.zeros((DIMENSIONS[1], DIMENSIONS[0], 4), dtype=numpy.uint8)
+
+    def replace(pixels: numpy.ndarray, position):
+        return replacement
+
+    rendered = list(frames.overlay_with(tracks, draw=replace))
+    for frame in rendered:
+        assert numpy.array_equal(frame.pixels, replacement)
+
+
+def test_overlay_raises_on_short_sidecar():
+    frames, _ = branches()
+    # Only one sidecar item, but the stream renders several frames.
+    video = frames.overlay_with([None], draw=draw_marker)
+    with pytest.raises(ValueError, match="not aligned"):
+        list(video)