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

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

import json
from concurrent.futures import ThreadPoolExecutor
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",
                 trabalhadores_em_paralelo: int = 4,
                 resolucao_maxima: int = 0) -> 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.")
        if resolucao_maxima < 0:
            raise ValueError("resolucao_maxima não pode ser negativa.")
        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
        self.trabalhadores_em_paralelo = max(1, trabalhadores_em_paralelo)
        # 0 preserva a resolução original; caso contrário, limita a maior dimensão
        # do frame extraído, sem nunca ampliar (force_original_aspect_ratio=decrease).
        self.resolucao_maxima = resolucao_maxima

    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)

        def extrair_um(indice_e_timestamp: tuple[int, float]) -> QuadroDeVideo:
            indice, timestamp = indice_e_timestamp
            # 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.
            # O nome carrega clipe, timestamp e resolução: dois clipes do mesmo arquivo
            # de origem não podem colidir no mesmo caminho, e mudar o intervalo ou a
            # resolução entre execuções não pode reaproveitar um frame de outra config.
            resolucao_rotulo = str(self.resolucao_maxima) if self.resolucao_maxima else "orig"
            destino = (self.diretorio_de_cache / arquivo.stem / clipe.identificador /
                       f"frame-{timestamp:09.3f}s-{resolucao_rotulo}.png")
            destino.parent.mkdir(parents=True, exist_ok=True)
            if destino.is_file():
                return QuadroDeVideo(timestamp, indice, largura, altura, None, destino)
            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}")
            return QuadroDeVideo(amostrado, indice, largura, altura, None, destino)

        tarefas = list(enumerate(timestamps))
        if self.trabalhadores_em_paralelo == 1 or len(tarefas) <= 1:
            return [extrair_um(tarefa) for tarefa in tarefas]
        with ThreadPoolExecutor(max_workers=self.trabalhadores_em_paralelo) as executor:
            return list(executor.map(extrair_um, tarefas))

    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."""
        comando_base = [self.ffmpeg, "-hide_banner", "-loglevel", "error"]
        filtro = (["-vf", f"scale=w={self.resolucao_maxima}:h={self.resolucao_maxima}:force_original_aspect_ratio=decrease"]
                  if self.resolucao_maxima else [])
        ultimo: subprocess.CompletedProcess[str] | None = None
        for candidato in (timestamp, max(inicio, timestamp - 0.1)):
            if destino.exists():
                destino.unlink()
            processo = subprocess.run([*comando_base, "-ss", str(candidato), "-i", str(arquivo),
                                        "-frames:v", "1", *filtro, "-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)