tools/video_encode_tool.py

# aduc_framework/tools/video_encode_tool.py
#
# Versão 1.4.0 (Conjunto de Ferramentas de Vídeo Finalizado)
# Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
#
# Este módulo atua como o especialista central para todas as operações de
# manipulação e codificação de vídeo. Ele abstrai as interações com
# FFmpeg e imageio, fornecendo uma API limpa e robusta para o resto do framework.
# - save_video_from_tensor: Converte um tensor de pixel em um arquivo .mp4.
# - extract_..._frame: Extrai frames específicos de clipes de vídeo.
# - concatenate_videos: Monta o filme final a partir dos clipes de cena.

import os
import subprocess
import logging
import random
import time
import shutil
from typing import List, Optional, Tuple

import imageio
import numpy as np
import torch

logger = logging.getLogger(__name__)

class VideoToolError(Exception):
    """Exceção personalizada para erros originados do VideoEncodeTool."""
    pass
                        

class VideoEncodeTool:
    def __init__(self, frame_log_every=8):
        self.frame_log_every = frame_log_every
        
    """
    Um especialista para lidar com tarefas de codificação e manipulação de vídeo.
    """


    @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}")

        pixel_5d = pixel_5d[:, :, -1:, :, :]
        
        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:
        """
        Extrai o primeiro frame de um arquivo de vídeo e o salva como uma imagem.
        """
        logger.info(f"Extraindo primeiro frame de '{os.path.basename(video_path)}'...")
        cmd = ['ffmpeg', '-y', '-v', 'error', '-i', video_path, '-vframes', '1', '-q:v', '2', output_image_path]
        try:
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            return output_image_path
        except subprocess.CalledProcessError as e:
            logger.error(f"FFmpeg (extract_first_frame) falhou: {e.stderr}")
            raise VideoToolError(f"Falha ao extrair o primeiro frame de {video_path}")

    def extract_last_frame(self, video_path: str, output_image_path: str) -> str:
        """
        Extrai o último frame de um arquivo de vídeo e o salva como uma imagem.
        """
        logger.info(f"Extraindo último frame de '{os.path.basename(video_path)}'...")
        cmd = ['ffmpeg', '-y', '-v', 'error', '-sseof', '-0.1', '-i', video_path, '-vframes', '1', '-q:v', '2', output_image_path]
        try:
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            return output_image_path
        except subprocess.CalledProcessError as e:
            logger.error(f"FFmpeg (extract_last_frame) falhou: {e.stderr}")
            raise VideoToolError(f"Falha ao extrair o último frame de {video_path}")

    def create_transition_bridge(self, start_image_path: str, end_image_path: str, 
                                 duration: float, fps: int, target_resolution: Tuple[int, int],
                                 workspace_dir: str, effect: Optional[str] = None) -> str:
        """
        Cria um clipe de vídeo curto que transiciona entre duas imagens estáticas.
        """
        output_path = os.path.join(workspace_dir, f"bridge_{int(time.time())}_{random.randint(100, 999)}.mp4")
        width, height = target_resolution
        fade_effects = ["fade", "wipeleft", "wiperight", "wipeup", "wipedown", "dissolve", "fadeblack", "fadewhite", "radial", "rectcrop", "circleopen", "circleclose", "horzopen", "horzclose"]
        selected_effect = effect if effect and effect.strip() else random.choice(fade_effects)
        transition_duration = max(0.1, duration)
        cmd = (f"ffmpeg -y -v error -loop 1 -t {transition_duration} -i \"{start_image_path}\" -loop 1 -t {transition_duration} -i \"{end_image_path}\" "
               f"-filter_complex \"[0:v]scale={width}:{height},setsar=1[v0];[1:v]scale={width}:{height},setsar=1[v1];"
               f"[v0][v1]xfade=transition={selected_effect}:duration={transition_duration}:offset=0[out]\" "
               f"-map \"[out]\" -c:v libx264 -r {fps} -pix_fmt yuv420p \"{output_path}\"")
        logger.info(f"Criando ponte de transição com efeito '{selected_effect}'...")
        try:
            subprocess.run(cmd, shell=True, check=True, text=True)
        except subprocess.CalledProcessError as e:
            raise VideoToolError(f"Falha ao criar vídeo de transição: {e.stderr}")
        return output_path

    def concatenate_videos(self, video_paths: List[str], output_path: str, workspace_dir: str) -> str:
        """
        Concatena múltiplos clipes de vídeo em um único arquivo sem re-codificar.
        """
        if not video_paths:
            raise VideoToolError("Nenhum fragmento de vídeo fornecido para concatenação.")
        # Se houver apenas um clipe, apenas o copie para o destino final.
        if len(video_paths) == 1:
            shutil.copy(video_paths[0], output_path)
            logger.info(f"Apenas um clipe fornecido. Copiado para '{output_path}'.")
            return output_path
            
        list_file_path = os.path.join(workspace_dir, f"concat_list_{int(time.time())}.txt")
        try:
            with open(list_file_path, 'w', encoding='utf-8') as f:
                for path in video_paths:
                    # Garante que o caminho seja absoluto para o ffmpeg encontrar
                    f.write(f"file '{os.path.abspath(path)}'\n")
            
            cmd_list = ['ffmpeg', '-y', '-f', 'concat', '-safe', '0', '-i', list_file_path, '-c', 'copy', output_path]
            logger.info(f"Concatenando {len(video_paths)} clipes para '{os.path.basename(output_path)}'...")
            subprocess.run(cmd_list, check=True, capture_output=True, text=True)
            logger.info("Concatenação FFmpeg bem-sucedida.")
            return output_path
        except subprocess.CalledProcessError as e:
            logger.error(f"Falha ao montar o vídeo final com FFmpeg: {e.stderr}")
            raise VideoToolError(f"Falha ao montar o vídeo final com FFmpeg.")
        finally:
            if os.path.exists(list_file_path):
                os.remove(list_file_path)

# --- Instância Singleton ---
video_encode_tool_singleton = VideoEncodeTool()