Create tools/video_encode_tool.py

tools/video_encode_tool.py

@@ -0,0 +1,41 @@
+# video_encode_tool.py — versão simples (beta 1.0)
+# Responsável por salvar um tensor 5D de pixels (B,C,T,H,W) em MP4, incrementalmente.
+
+import torch
+import imageio
+import numpy as np
+
+class _SimpleVideoEncoder:
+    def __init__(self, frame_log_every=8):
+        self.frame_log_every = frame_log_every
+
+    @torch.no_grad()
+    def save_video_from_tensor(self, pixel_5d: torch.Tensor, path: str, fps: int = 24, progress_callback=None):
+        """
+        Espera pixel_5d em [0,1], shape (B,C,T,H,W).
+        Escreve MP4 incremental, convertendo cada frame para (H,W,C) uint8.
+        """
+        # Move para CPU apenas para formar os frames HWC uint8 com baixo overhead
+        device = "cuda" if pixel_5d.is_cuda else "cpu"
+        B, C, T, H, W = pixel_5d.shape
+        if B != 1:
+            # Mantemos simples: um vídeo por chamada (B=1)
+            raise ValueError(f"Esperado B=1, recebido B={B}")
+
+        with imageio.get_writer(path, fps=int(fps), codec="libx264", quality=8) as writer:
+            for i in range(T):
+                frame_chw = pixel_5d[0, :, i]  # (C,H,W)
+                frame_hwc_u8 = (frame_chw.permute(1, 2, 0)
+                                .clamp(0, 1)
+                                .mul(255)
+                                .to(torch.uint8)
+                                .cpu()
+                                .numpy())
+                writer.append_data(frame_hwc_u8)
+                if progress_callback:
+                    progress_callback(i + 1, T)
+                if i % self.frame_log_every == 0:
+                    print(f"[DEBUG] [Encoder] frame {i}/{T} gravado ({H}x{W}@{fps}fps)")
+
+# Singleton global de uso simples
+video_encode_tool_singleton = _SimpleVideoEncoder()