diff --git a/examples/blob_tracker_overlay.py b/examples/blob_tracker_overlay.py new file mode 100644 index 0000000..0ce60b2 --- /dev/null +++ b/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() diff --git a/python/faery/frame_filter.py b/python/faery/frame_filter.py index f4bcbb1..c541a17 100644 --- a/python/faery/frame_filter.py +++ b/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 diff --git a/python/faery/frame_stream.py b/python/faery/frame_stream.py index 635d158..8a28feb 100644 --- a/python/faery/frame_stream.py +++ b/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"): diff --git a/tests/test_overlay.py b/tests/test_overlay.py new file mode 100644 index 0000000..90898ab --- /dev/null +++ b/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)