tools/video_encode_tool.py

# 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-1):
                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()