diff --git a/positronic/offboard/frame_codec_cost.py b/positronic/offboard/frame_codec_cost.py
new file mode 100644
index 000000000..21350711f
--- /dev/null
+++ b/positronic/offboard/frame_codec_cost.py
@@ -0,0 +1,175 @@
+"""Measure what one observation window costs per image codec: bytes, encode time, decode time.
+
+Replays a recorded episode's cameras through the rig-side bound, then encodes each temporal-stack
+window as one JPEG per frame, which the wire carries, and as one h264 GOP. JPEG encodes
+single-threaded through ``encode_jpeg``; h264 encodes with x264's own frame threading.
+
+Usage
+ python -m positronic.offboard.frame_codec_cost --episode
.mp4>
+ python -m positronic.offboard.frame_codec_cost --episode --bound 640x180 --json rows.json
+"""
+
+import argparse
+import dataclasses
+import io
+import json
+import pathlib
+import statistics
+import time
+
+import av
+import numpy as np
+
+from positronic.policy.codec import RestrictImageSize
+from positronic.utils.serialization import FRAMES, encode_jpeg, unpack
+
+
+@dataclasses.dataclass(frozen=True)
+class Jpeg:
+ """One JPEG per frame, at the quality the wire itself encodes at."""
+
+ name = 'jpeg'
+
+ def cost(self, window: np.ndarray) -> tuple[int, float, float]:
+ """Bytes, encode ms and decode ms for one window."""
+ start = time.perf_counter()
+ marker = encode_jpeg(window)
+ encode_ms = 1000 * (time.perf_counter() - start)
+
+ start = time.perf_counter()
+ unpack(marker)
+ decode_ms = 1000 * (time.perf_counter() - start)
+ return sum(len(buf) for buf in marker[FRAMES]), encode_ms, decode_ms
+
+
+@dataclasses.dataclass(frozen=True)
+class H264:
+ """One x264 setting, over the whole window as a single GOP."""
+
+ preset: str
+ crf: int
+
+ @property
+ def name(self) -> str:
+ return f'h264 {self.preset} crf{self.crf}'
+
+ def cost(self, window: np.ndarray) -> tuple[int, float, float]:
+ """Bytes, encode ms and decode ms for one window."""
+ height, width = window.shape[1:3]
+ buffer = io.BytesIO()
+ start = time.perf_counter()
+ with av.open(buffer, 'w', format='mp4') as container:
+ stream = container.add_stream('libx264', rate=15)
+ stream.width, stream.height, stream.pix_fmt = width, height, 'yuv420p'
+ # One self-contained GOP with no lookahead: a request carries its own window and waits on it.
+ stream.options = {'preset': self.preset, 'crf': str(self.crf), 'tune': 'zerolatency', 'g': str(len(window))}
+ for frame in window:
+ container.mux(stream.encode(av.VideoFrame.from_ndarray(frame, format='rgb24')))
+ container.mux(stream.encode(None))
+ encode_ms = 1000 * (time.perf_counter() - start)
+ payload = buffer.getvalue()
+
+ start = time.perf_counter()
+ with av.open(io.BytesIO(payload), 'r') as container:
+ for frame in container.decode(video=0):
+ frame.to_ndarray(format='rgb24')
+ return len(payload), encode_ms, 1000 * (time.perf_counter() - start)
+
+
+Codec = Jpeg | H264
+
+
+@dataclasses.dataclass(frozen=True)
+class Cost:
+ """What one window of one camera cost under one codec."""
+
+ codec: str
+ camera: str
+ window: int
+ kib: float
+ encode_ms: float
+ decode_ms: float
+
+
+def bounded_frames(mp4: pathlib.Path, bound: RestrictImageSize, rate_hz: float) -> list[np.ndarray]:
+ """Every frame the stack samples, through the rig's own bound."""
+ key = 'image'
+ with av.open(str(mp4), 'r') as container:
+ recorded_rate = container.streams.video[0].average_rate
+ if recorded_rate is None:
+ raise ValueError(f'{mp4} declares no frame rate, so the sampled frames cannot be chosen')
+ step = max(1, round(float(recorded_rate) / rate_hz))
+ frames = []
+ for index, frame in enumerate(container.decode(video=0)):
+ if index % step:
+ continue
+ frames.append(bound.encode({key: frame.to_ndarray(format='rgb24')})[key])
+ return frames
+
+
+def costs(frames: list[np.ndarray], camera: str, depth: int, codecs: list[Codec], limit: int) -> list[Cost]:
+ """One row per codec per window, over the windows the episode holds."""
+ starts = list(range(0, len(frames) - depth + 1))[: limit or None]
+ rows = []
+ for window_index, start in enumerate(starts):
+ window = np.stack(frames[start : start + depth])
+ for codec in codecs:
+ size, encode_ms, decode_ms = codec.cost(window)
+ rows.append(Cost(codec.name, camera, window_index, size / 1024, encode_ms, decode_ms))
+ return rows
+
+
+def _row(label: str, group: list[Cost], windows: list[int]) -> str:
+ """One printed line: the median over windows of what the group cost in each window."""
+ per_window = [[row for row in group if row.window == window] for window in windows]
+ return (
+ f'{label:>16} '
+ f'{statistics.median(sum(row.kib for row in w) for w in per_window):8.0f} '
+ f'{statistics.median(sum(row.encode_ms for row in w) for w in per_window):10.1f} '
+ f'{statistics.median(sum(row.decode_ms for row in w) for w in per_window):10.1f}'
+ )
+
+
+def report(rows: list[Cost], depth: int) -> None:
+ """Median cost per codec, per camera and summed over the cameras one request carries."""
+ names = list(dict.fromkeys(row.codec for row in rows))
+ cameras = list(dict.fromkeys(row.camera for row in rows))
+ windows = sorted({row.window for row in rows})
+ print(f'\n{depth} frames, {len(windows)} windows, {len(cameras)} cameras\n')
+ print(f'{"codec":>22} {"camera":>16} {"KiB":>8} {"encode ms":>10} {"decode ms":>10}')
+ for name in names:
+ of_codec = [row for row in rows if row.codec == name]
+ for camera in cameras:
+ print(f'{name:>22} {_row(camera, [row for row in of_codec if row.camera == camera], windows)}')
+ print(f'{name:>22} {_row("every camera", of_codec, windows)}')
+
+
+def main() -> int:
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--episode', required=True, type=pathlib.Path, help='directory holding .mp4')
+ parser.add_argument('--cameras', default='image.exterior,image.wrist')
+ parser.add_argument('--frames', type=int, default=25, help='temporal-stack depth')
+ parser.add_argument('--rate', type=float, default=15.0, help='stack sampling rate, Hz')
+ parser.add_argument('--bound', default='1024x288', help='WxH rig-side bound')
+ parser.add_argument('--x264', default='ultrafast:20,veryfast:20', help='comma-separated preset:crf')
+ parser.add_argument('--windows', type=int, default=100, help='windows per camera, 0 for every one')
+ parser.add_argument('--json', type=pathlib.Path, help='write every row here')
+ args = parser.parse_args()
+
+ bound = RestrictImageSize(*(int(side) for side in args.bound.lower().split('x')))
+ codecs: list[Codec] = [Jpeg()]
+ codecs += [H264(spec.split(':')[0], int(spec.split(':')[1])) for spec in args.x264.split(',')]
+ rows: list[Cost] = []
+ for camera in args.cameras.split(','):
+ frames = bounded_frames(args.episode / f'{camera}.mp4', bound, args.rate)
+ print(f'{camera}: {len(frames)} sampled frames at {frames[0].shape[1]}x{frames[0].shape[0]}')
+ rows += costs(frames, camera, args.frames, codecs, args.windows)
+
+ report(rows, args.frames)
+ if args.json:
+ args.json.write_text(json.dumps([dataclasses.asdict(row) for row in rows]))
+ return 0
+
+
+if __name__ == '__main__':
+ raise SystemExit(main())
diff --git a/positronic/utils/serialization.py b/positronic/utils/serialization.py
index 09948c22a..e7552dd8f 100644
--- a/positronic/utils/serialization.py
+++ b/positronic/utils/serialization.py
@@ -26,7 +26,7 @@
_DATA = b'data'
_DTYPE = b'dtype'
_SHAPE = b'shape'
-_FRAMES = b'frames'
+FRAMES = b'frames' # the wire's own name for the per-frame JPEGs
_NDIM = b'ndim'
# JPEG quality for images on the wire. A single HD frame — and especially a (T, H, W, 3) stack — is many
@@ -46,12 +46,12 @@ def encode_jpeg(image: np.ndarray) -> dict[bytes, Any]:
buf = io.BytesIO()
PilImage.fromarray(np.ascontiguousarray(frame, dtype=np.uint8)).save(buf, format='JPEG', quality=_JPEG_QUALITY)
bufs.append(buf.getvalue())
- return {_JPEG: True, _FRAMES: bufs, _NDIM: int(image.ndim)}
+ return {_JPEG: True, FRAMES: bufs, _NDIM: int(image.ndim)}
def _decode_jpeg(marker: dict) -> np.ndarray:
"""Inverse of ``encode_jpeg``: decode per-frame JPEGs and restore the original shape."""
- frames = np.stack([np.asarray(PilImage.open(io.BytesIO(buf))) for buf in marker[_FRAMES]])
+ frames = np.stack([np.asarray(PilImage.open(io.BytesIO(buf))) for buf in marker[FRAMES]])
return frames if marker[_NDIM] == 4 else frames[0]