eval: log validation videos to wandb

egomimic/eval/eval_hpt.py

@@ -126,10 +126,3 @@ def compute_metrics_and_viz(self, batch, do_viz=True):
             metrics["Valid/action_loss"] = total_loss / n_loss_embodiments
 
         return metrics, images_dict
-
-    def _visualize_preds(self, predictions, batch):
-        if self.viz_func is None:
-            raise ValueError("viz_func is not set")
-        embodiment_id = batch["embodiment"][0].item()
-        embodiment_name = get_embodiment(embodiment_id).lower()
-        return self.viz_func[embodiment_name](predictions, batch)

egomimic/eval/eval_latent.py

@@ -61,7 +61,8 @@ def __init__(
         emit_combined: bool = True,
         color_by: str = "embodiment",  # "embodiment" or "hash"
     ):
-        super().__init__(limit_val_batches=limit_val_batches)
+        # Its own _visualize_preds renders overlay frames only.
+        super().__init__(limit_val_batches=limit_val_batches, viz_mode="overlay")
         self.compute_umap = compute_umap
         self.compute_tsne_2d = compute_tsne_2d
         self.compute_tsne_3d = compute_tsne_3d

egomimic/eval/eval_pi.py

@@ -90,10 +90,3 @@ def compute_metrics_and_viz(self, batch, do_viz=True):
             metrics["Valid/action_loss"] = total_loss / n_loss_embodiments
 
         return metrics, images_dict
-
-    def _visualize_preds(self, predictions, batch):
-        if self.viz_func is None:
-            raise ValueError("viz_func is not set")
-        embodiment_id = batch["embodiment"][0].item()
-        embodiment_name = get_embodiment(embodiment_id).lower()
-        return self.viz_func[embodiment_name](predictions, batch)

egomimic/eval/eval_train_viz.py

@@ -35,8 +35,17 @@ def __init__(self, base: EvalVideo, prefix: str = "train_viz"):
             transform_lists=base.transform_lists,
             viz_every_n_epochs=base.viz_every_n_epochs,
             viz_max_batches=base.viz_max_batches,
+            viz_mode=base.viz_mode,
+            replay_chunks=base.replay_chunks,
+            replay_trail=base.replay_trail,
+            action_stride=base.action_stride,
         )
 
+    def _set_replay_now(self, value: bool) -> None:
+        # The base renders the batch, so it must know which mode this step is.
+        super()._set_replay_now(value)
+        self.base._set_replay_now(value)
+
     @property
     def trainer(self):
         return self._trainer

egomimic/eval/eval_video.py

@@ -3,14 +3,34 @@
 
 import torch
 import torchvision.io as tvio
+from lightning.pytorch.loggers import WandbLogger
 
 from egomimic.eval.eval import Eval
+from egomimic.eval.replay_viz import (
+    MIN_FRAMES_FOR_STRIDE,
+    ReplayClip,
+    infer_stride,
+    render_replay,
+)
 from egomimic.rldb.embodiment.embodiment import get_embodiment
 
 # The source rate: val loaders are never sharded across ranks, so rank 0
 # renders every source frame.
 VIDEO_FPS = 30
 
+# viz_gt_preds kwargs that ReplayClip / render_replay consume themselves.
+_REPLAY_OWN_KWARGS = {
+    "image_key",
+    "action_key",
+    "annotation_key",
+    "mode",
+    "gt_alpha",
+    "pred_alpha",
+}
+
+# Flush a key's buffer to a video part past this many frames (host memory).
+_MAX_BUFFERED_FRAMES = 1000
+
 
 class EvalVideo(Eval):
     """
@@ -26,6 +46,10 @@ def __init__(
         transform_lists: dict | None = None,
         viz_every_n_epochs: int = 1,
         viz_max_batches: int | None = None,
+        viz_mode: str = "replay",
+        replay_chunks: int = 5,
+        replay_trail: int = 5,
+        action_stride: int | None = None,
     ):
         super().__init__()
         self.trainer = None
@@ -39,6 +63,26 @@ def __init__(
         # wall time; metrics are still computed on EVERY batch
         # regardless. Required on heads with pinned video loaders (trainHydra).
         self.viz_max_batches = viz_max_batches
+        # "replay" (egomimic/eval/replay_viz.py) plays each chunk as GT then
+        # prediction; it needs consecutive frames, so it runs only on a pinned
+        # video loader and caps it at ``replay_chunks`` chunks instead of
+        # ``viz_max_batches``. "overlay" draws each frame's whole GT and
+        # predicted chunk, and is what a head without a video loader gets.
+        # ``action_stride`` is frames per chunk step (None: inferred from the
+        # GT keypoints, see replay_viz.infer_stride). The cap counts frames
+        # across a head's pinned episodes in config order, so a later episode
+        # plays only when the earlier ones are shorter than the cap; chunks
+        # never straddle two episodes.
+        if viz_mode not in ("overlay", "replay"):
+            raise ValueError(f"unknown viz_mode {viz_mode!r}")
+        self.viz_mode = viz_mode
+        self.replay_chunks = replay_chunks
+        self._replay_now = False
+        self._replay_seen = {}
+        self._replay_stride = {}
+        self._replay_horizon = {}
+        self.replay_trail = replay_trail
+        self.action_stride = action_stride
         # Per-embodiment list[Transform] applied once during eval to project
         # the model's wrist-frame actions back into cam (head) frame for the
         # viz video.
@@ -93,6 +137,67 @@ def compute_metrics_and_viz(self, batch, do_viz=True):
         """
         raise NotImplementedError
 
+    def _visualize_preds(self, predictions, batch):
+        if self.viz_func is None:
+            raise ValueError("viz_func is not set")
+        embodiment_name = get_embodiment(batch["embodiment"][0].item()).lower()
+        viz_fn = self.viz_func[embodiment_name]
+        if self._replay_now:
+            return ReplayClip.from_batch(viz_fn, predictions, batch)
+        return viz_fn(predictions, batch)
+
+    def _render_replay(self, key, final: bool):
+        """Render the buffered clip's complete chunks; keep the rest unless final."""
+        clips = self.val_image_buffer[key]
+        if not clips:
+            return None
+        clip = ReplayClip.concat(clips)
+        viz_fn = self.viz_func[get_embodiment(key).lower()]
+        frames, used = render_replay(
+            clip,
+            embodiment_cls=viz_fn.func.__self__,
+            mode=viz_fn.keywords["mode"],
+            stride=self._replay_stride.get(key, self.action_stride),
+            trail=self.replay_trail,
+            viz_kwargs={
+                k: v for k, v in viz_fn.keywords.items() if k not in _REPLAY_OWN_KWARGS
+            },
+        )
+        self.val_image_buffer[key] = [] if final else [clip.tail(used)]
+        return torch.from_numpy(frames) if len(frames) else None
+
+    def _set_replay_now(self, value: bool) -> None:
+        self._replay_now = value
+
+    def _is_replay_buffer(self, key) -> bool:
+        buffer = self.val_image_buffer[key]
+        return bool(buffer) and isinstance(buffer[0], ReplayClip)
+
+    def _replay_done(self) -> bool:
+        return bool(self._replay_seen) and all(
+            k in self._replay_stride
+            and seen
+            >= self.replay_chunks * self._replay_horizon[k] * self._replay_stride[k]
+            for k, seen in self._replay_seen.items()
+        )
+
+    def _buffered_frames(self, key) -> int:
+        # Overlay buffers hold one (H, W, 3) tensor per frame, not clips.
+        if self._is_replay_buffer(key):
+            return sum(len(c) for c in self.val_image_buffer[key])
+        return len(self.val_image_buffer[key])
+
+    def _flush(self, key, final: bool) -> None:
+        if self._is_replay_buffer(key):
+            frames = self._render_replay(key, final)
+        else:
+            buffer = self.val_image_buffer[key]
+            frames = torch.stack(buffer) if buffer else None
+            self.val_image_buffer[key] = []
+        if frames is not None:
+            self._write_video(key, frames)
+            self.val_counter[key] += 1
+
     def _write_video(self, key, frames) -> None:
         # Every rank that gets here renders the same frames to the same path
         # (no DistributedSampler splits the val loaders), so without
@@ -107,6 +212,21 @@ def _write_video(self, key, frames) -> None:
         )
         os.makedirs(os.path.dirname(path), exist_ok=True)
         tvio.write_video(path, frames, fps=VIDEO_FPS, video_codec="h264")
+        self._upload_video(key, path)
+
+    def _upload_video(self, key, path) -> None:
+        name = f"{os.path.basename(self.video_dir())}/{get_embodiment(key)}"
+        if self.val_counter[key]:
+            name += f"_{self.val_counter[key]}"
+        for logger in self.trainer.loggers:
+            if isinstance(logger, WandbLogger):
+                logger.log_video(
+                    name,
+                    [path],
+                    step=self.trainer.global_step,
+                    caption=[f"epoch {self.trainer.current_epoch}"],
+                    format=["mp4"],
+                )
 
     def on_validation_start(self):
         # _write_video creates the epoch directory it writes into.
@@ -115,11 +235,11 @@ def on_validation_start(self):
     def on_validation_end(self):
         if not self._should_viz():
             return
-        for key, buffer in self.val_image_buffer.items():
-            if len(buffer) != 0:
-                self._write_video(key, torch.stack(buffer))
+        for key in list(self.val_image_buffer):
+            self._flush(key, final=True)
             self.val_counter[key] = 0
             self.val_image_buffer[key] = []
+        self._replay_seen, self._replay_stride = {}, {}
 
     def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"):
         """``mode`` splits metrics from video so one head can serve two loaders:
@@ -129,16 +249,26 @@ def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"):
           subsampled video would be an incoherent time-lapse anyway).
         * ``"video"`` -- the contiguous pinned-episode loader. Renders and
           buffers, and logs NOTHING, so no ``Valid/...`` key is averaged over
-          it. On a non-viz epoch, or past ``viz_max_batches``, it returns before
-          the forward pass: there is nothing such a call could produce.
+          it. On a non-viz epoch, or past its cap (``viz_max_batches``, or
+          ``replay_chunks`` chunks in replay mode), it returns before the
+          forward pass: there is nothing such a call could produce.
         * ``"both"`` (default) -- today's single-loader behaviour, which is what
           a data config without ``video_episodes`` still gets.
         """
         if mode not in ("both", "metrics", "video"):
             raise ValueError(f"unknown validation mode {mode!r}")
-        past_viz_cap = (
-            self.viz_max_batches is not None and batch_idx >= self.viz_max_batches
-        )
+        replay = mode == "video" and self.viz_mode == "replay"
+        self._set_replay_now(replay)
+        if replay:
+            # viz_max_batches still bounds a loader that yields no clips.
+            nothing_yet = not self._replay_seen and (
+                self.viz_max_batches is not None and batch_idx >= self.viz_max_batches
+            )
+            past_viz_cap = self._replay_done() or nothing_yet
+        else:
+            past_viz_cap = (
+                self.viz_max_batches is not None and batch_idx >= self.viz_max_batches
+            )
         if mode == "video" and (past_viz_cap or not self._should_viz()):
             return
         do_viz = mode != "metrics" and self._should_viz() and not past_viz_cap
@@ -158,11 +288,22 @@ def on_validation_step(self, batch, batch_idx, dataloader_idx=0, mode="both"):
                 ):
                     self.val_image_buffer[key] = []
                     self.val_counter[key] = 0
-                self.val_image_buffer[key].extend(torch.from_numpy(images))
-                if len(self.val_image_buffer[key]) >= 1000:
-                    self._write_video(key, torch.stack(self.val_image_buffer[key]))
-                    self.val_image_buffer[key].clear()
-                    self.val_counter[key] += 1
+                if isinstance(images, ReplayClip):
+                    self.val_image_buffer[key].append(images)
+                    seen = self._replay_seen.get(key, 0) + len(images)
+                    self._replay_seen[key] = seen
+                    self._replay_horizon[key] = images.gt.shape[1]
+                    if key not in self._replay_stride and seen >= (
+                        MIN_FRAMES_FOR_STRIDE
+                    ):
+                        buffered = ReplayClip.concat(self.val_image_buffer[key])
+                        self._replay_stride[key] = (
+                            self.action_stride or infer_stride(buffered) or 1
+                        )
+                else:
+                    self.val_image_buffer[key].extend(torch.from_numpy(images))
+                if self._buffered_frames(key) >= _MAX_BUFFERED_FRAMES:
+                    self._flush(key, final=False)
 
         if mode == "video":
             # Video-only loader: the forward pass was for the overlay. Logging