Update tools/video_encode_tool.py

tools/video_encode_tool.py

@@ -27,51 +27,44 @@ logger = logging.getLogger(__name__)
 class VideoToolError(Exception):
     """Exceção personalizada para erros originados do VideoEncodeTool."""
     pass
+                        
 
 class VideoEncodeTool:
     """
     Um especialista para lidar com tarefas de codificação e manipulação de vídeo.
     """
 
-    def save_video_from_tensor(self, video_tensor: torch.Tensor, path: str, fps: int = 24):
+
+    @torch.no_grad()
+    def save_video_from_tensor(self, pixel_5d: torch.Tensor, path: str, fps: int = 24, progress_callback=None):
         """
-        Salva um tensor de pixel como um arquivo de vídeo .mp4 usando parâmetros otimizados.
-        Espera um tensor no formato (B, C, F, H, W) onde B=1.
+        Espera pixel_5d em [0,1], shape (B,C,T,H,W).
+        Escreve MP4 incremental, convertendo cada frame para (H,W,C) uint8.
         """
-        # Verificações de robustez para garantir que o tensor é válido
-        if video_tensor is None or video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[2] == 0:
-            logger.warning(f"Tensor de vídeo inválido ou vazio recebido. Shape: {video_tensor.shape if video_tensor is not None else 'None'}. Pulando salvamento de vídeo para '{path}'.")
-            return
+        # 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}")
 
-        logger.info(f"Salvando tensor de vídeo com shape {video_tensor.shape} para '{os.path.basename(path)}'...")
+        pixel_5d = pixel_5d[:, :, :-1, :, :]
         
-        try:
-            # Squeeze: (1, C, F, H, W) -> (C, F, H, W)
-            # Permute: (C, F, H, W) -> (F, H, W, C) - formato esperado por imageio
-            video_tensor_permuted = video_tensor.squeeze(0).permute(1, 2, 3, 0)
-            
-            # Desnormaliza de [-1, 1] para [0, 1]
-            video_tensor_normalized = (video_tensor_permuted.clamp(-1, 1) + 1) / 2.0
-            
-            # Converte para [0, 255], move para CPU e converte para numpy uint8
-            video_np = (video_tensor_normalized.detach().cpu().float().numpy() * 255).astype(np.uint8)
-            
-            # Salva o vídeo com parâmetros de alta compatibilidade
-            with imageio.get_writer(
-                path, 
-                fps=fps, 
-                codec='libx264', 
-                quality=8, # Qualidade boa (0-10, onde 10 é a melhor)
-                output_params=['-pix_fmt', 'yuv420p'] # Formato de pixel para compatibilidade máxima
-            ) as writer:
-                for frame in video_np:
-                    writer.append_data(frame)
-                     
-            
-            logger.info(f"Vídeo salvo com sucesso em: {path}")
-        except Exception as e:
-            logger.error(f"Falha ao salvar vídeo com imageio para '{path}': {e}", exc_info=True)
-            raise VideoToolError(f"Não foi possível escrever o arquivo de vídeo: {e}")
+        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)")
+                    
 
     def extract_first_frame(self, video_path: str, output_image_path: str) -> str:
         """