| """Video recording with the per-frame HUD overlay. |
| |
| Also owns the warm-up-frame drop: the first frames after a reset render with the wrong camera |
| pose and unresolved textures, which made reviewers describe scenes (missing containers) that |
| never existed in the episode. |
| """ |
| from __future__ import annotations |
|
|
| import os |
|
|
| import numpy as np |
|
|
| WARMUP_FRAMES = 2 |
|
|
|
|
| class Recorder: |
| def __init__(self, env, title="", episode=-1, fps=14): |
| self.env = env |
| self.title = title |
| self.episode = episode |
| self.fps = fps |
| self.frames = [] |
| self.phase = "start" |
| self.result = "" |
| self.lines_fn = None |
|
|
| def capture(self): |
| img = self.env.render() |
| if img is None: |
| return |
| from PIL import Image, ImageDraw |
| im = Image.fromarray(np.asarray(img)[..., :3].copy()) |
| d = ImageDraw.Draw(im) |
| lines = [f"=== {self.title} ===" + (f" EP {self.episode}" if self.episode >= 0 else "")] |
| if self.result: |
| lines.append(f"RESULT: {self.result}") |
| lines.append(f"ACTION: {self.phase}") |
| if self.lines_fn is not None: |
| try: |
| lines += list(self.lines_fn()) |
| except Exception: |
| pass |
| d.rectangle([0, 0, 470, 18*len(lines)+6], fill=(0, 0, 0)) |
| y = 3 |
| for ln in lines: |
| d.text((6, y), ln, fill=(255, 235, 60)); y += 18 |
| self.frames.append(np.array(im)) |
|
|
| def save(self, path): |
| import imageio.v2 as imageio |
| frames = self.frames[WARMUP_FRAMES:] if len(self.frames) > WARMUP_FRAMES+4 else self.frames |
| if not frames: |
| print("[rec] no frames to write"); return None |
| os.makedirs(os.path.dirname(path) or ".", exist_ok=True) |
| imageio.mimsave(path, frames, fps=self.fps) |
| print(f"[rec] video -> {path} ({len(frames)} frames)", flush=True) |
| return path |
|
|
