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

from pathlib import Path
from typing import Any

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


class ExtratorDeQuadrosOpenCV(ExtratorDeQuadros):
    """Extrai frames em intervalos regulares sem expor detalhes do OpenCV."""

    def __init__(self, intervalo_em_segundos: float = 1.0,
                 diretorio_de_cache: str | Path | None = None,
                 maximo_de_quadros: int | None = None,
                 cv2_module: Any | None = None) -> 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) if diretorio_de_cache else None
        self.maximo_de_quadros = maximo_de_quadros
        self._cv2 = cv2_module

    @property
    def cv2(self) -> Any:
        if self._cv2 is None:
            try:
                import cv2
            except ImportError as exc:
                raise RuntimeError("OpenCV não está instalado. Instale opencv-python.") from exc
            self._cv2 = cv2
        return self._cv2

    def extrair(self, clipe: Clipe) -> list[QuadroDeVideo]:
        if not clipe.arquivo:
            raise ValueError("Clipe não possui arquivo de origem.")
        caminho = Path(clipe.arquivo)
        if not caminho.is_file():
            raise FileNotFoundError(f"Arquivo do clipe não encontrado: {caminho}")
        captura = self.cv2.VideoCapture(str(caminho))
        if not captura.isOpened():
            raise RuntimeError(f"OpenCV não conseguiu abrir o vídeo: {caminho}")
        try:
            fps = float(captura.get(self.cv2.CAP_PROP_FPS) or 0)
            total_de_frames = int(captura.get(self.cv2.CAP_PROP_FRAME_COUNT) or 0)
            duracao = total_de_frames / fps if fps > 0 else 0.0
            inicio, fim = self._intervalo_de_origem(clipe, duracao)
            timestamps = self._timestamps(inicio, fim)
            resultado: list[QuadroDeVideo] = []
            for indice, timestamp in enumerate(timestamps):
                captura.set(self.cv2.CAP_PROP_POS_MSEC, timestamp * 1000.0)
                sucesso, imagem = captura.read()
                if not sucesso or imagem is None:
                    continue
                altura, largura = imagem.shape[:2]
                salvo = self._salvar(caminho, indice, imagem)
                resultado.append(QuadroDeVideo(timestamp, indice, largura, altura, imagem, salvo))
            return resultado
        finally:
            captura.release()

    def _timestamps(self, inicio: float, fim: float) -> list[float]:
        if fim <= inicio:
            return [inicio]
        quantidade = int((fim - inicio) / self.intervalo_em_segundos) + 1
        limite = max(inicio, fim - 0.001)
        timestamps = [min(inicio + indice * self.intervalo_em_segundos, limite)
                      for indice in range(quantidade)]
        return timestamps[:self.maximo_de_quadros]

    @staticmethod
    def _intervalo_de_origem(clipe: Clipe, duracao: float) -> tuple[float, float]:
        if clipe.intervalo_na_origem is None:
            return 0.0, duracao
        inicio = max(0.0, clipe.intervalo_na_origem.inicio)
        fim = min(duracao, clipe.intervalo_na_origem.fim) if duracao > 0 else clipe.intervalo_na_origem.fim
        return inicio, max(inicio, fim)

    def _salvar(self, arquivo: Path, indice: int, imagem: Any) -> Path | None:
        if self.diretorio_de_cache is None:
            return None
        destino = self.diretorio_de_cache / arquivo.stem / f"frame-{indice:06d}.jpg"
        destino.parent.mkdir(parents=True, exist_ok=True)
        if not self.cv2.imwrite(str(destino), imagem):
            raise RuntimeError(f"Não foi possível salvar frame: {destino}")
        return destino