File size: 1,730 Bytes
db005a9 e560c23 db005a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
# 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()
|