jhonny-editor/code/engine/integracoes/visual/extrator_ffmpeg.py

"""Extrator de frames para Vision usando somente ferramentas locais do sistema."""

import json
from pathlib import Path
import subprocess

from ...dominio import Clipe
from .contratos import ExtratorDeQuadros
from .extrator_open_cv import ExtratorDeQuadrosOpenCV
from .modelos import QuadroDeVideo


class ExtratorDeQuadrosFFmpeg(ExtratorDeQuadros):
    """Persiste amostras de um arquivo original sem carregar OpenCV no processo."""

    def __init__(self, intervalo_em_segundos: float = 1.0,
                 diretorio_de_cache: str | Path = ".frames",
                 maximo_de_quadros: int | None = None,
                 ffmpeg: str = "ffmpeg", ffprobe: str = "ffprobe") -> None:
        if intervalo_em_segundos <= 0:
            raise ValueError("intervalo_em_segundos deve ser positivo.")
        if maximo_de_quadros is not None and maximo_de_quadros < 1:
            raise ValueError("maximo_de_quadros deve ser maior que zero.")
        self.intervalo_em_segundos = intervalo_em_segundos
        self.diretorio_de_cache = Path(diretorio_de_cache)
        self.maximo_de_quadros = maximo_de_quadros
        self.ffmpeg = ffmpeg
        self.ffprobe = ffprobe

    def extrair(self, clipe: Clipe) -> list[QuadroDeVideo]:
        if not clipe.arquivo:
            raise ValueError("Clipe não possui arquivo de origem.")
        arquivo = Path(clipe.arquivo)
        if not arquivo.is_file():
            raise FileNotFoundError(f"Arquivo do clipe não encontrado: {arquivo}")
        largura, altura, duracao = self._metadados(arquivo)
        inicio, fim = ExtratorDeQuadrosOpenCV._intervalo_de_origem(clipe, duracao)
        timestamps = ExtratorDeQuadrosOpenCV(self.intervalo_em_segundos,
                                              maximo_de_quadros=self.maximo_de_quadros)._timestamps(inicio, fim)
        resultado: list[QuadroDeVideo] = []
        for indice, timestamp in enumerate(timestamps):
            # PNG evita a recodificação JPEG de YUV limitado, que falha em parte
            # dos vídeos móveis e ainda preserva melhor OCR/segmentação para Vision.
            destino = self.diretorio_de_cache / arquivo.stem / f"frame-{indice:06d}.png"
            destino.parent.mkdir(parents=True, exist_ok=True)
            amostrado, processo = self._extrair_frame(arquivo, destino, timestamp, inicio)
            if processo.returncode != 0 or not destino.is_file():
                detalhe = processo.stderr.strip() or "ffmpeg não produziu o frame"
                raise RuntimeError(f"Não foi possível extrair frame em {timestamp}s: {detalhe}")
            resultado.append(QuadroDeVideo(amostrado, indice, largura, altura, None, destino))
        return resultado

    def _extrair_frame(self, arquivo: Path, destino: Path, timestamp: float,
                        inicio: float) -> tuple[float, subprocess.CompletedProcess[str]]:
        """Recua um pouco no fim do arquivo se o decoder não encontrar frame no timestamp."""
        ultimo: subprocess.CompletedProcess[str] | None = None
        for candidato in (timestamp, max(inicio, timestamp - 0.1)):
            if destino.exists():
                destino.unlink()
            processo = subprocess.run([self.ffmpeg, "-hide_banner", "-loglevel", "error", "-ss", str(candidato),
                                        "-i", str(arquivo), "-frames:v", "1", "-y", str(destino)],
                                       capture_output=True, text=True, timeout=60)
            if processo.returncode == 0 and destino.is_file():
                return candidato, processo
            ultimo = processo
        assert ultimo is not None
        return timestamp, ultimo

    def _metadados(self, arquivo: Path) -> tuple[int, int, float]:
        processo = subprocess.run([self.ffprobe, "-v", "error", "-select_streams", "v:0",
                                   "-show_entries", "stream=width,height:format=duration",
                                   "-of", "json", str(arquivo)], capture_output=True, text=True, timeout=30)
        if processo.returncode != 0:
            raise RuntimeError(processo.stderr.strip() or "ffprobe não conseguiu ler o vídeo")
        dados = json.loads(processo.stdout)
        stream = (dados.get("streams") or [{}])[0]
        return int(stream["width"]), int(stream["height"]), float((dados.get("format") or {}).get("duration") or 0)