THOR / ffmpeg.py
ElJoker63's picture
Actualizar nodo OmniVision
3216f7c verified
Raw History Blame Contribute Delete
18.6 kB
"""
Módulo de procesamiento multimedia con FFmpeg y sondeo con ffprobe.
Port 1 a 1 de internal/ffmpeg/ffmpeg.go con todas las validaciones y sanitizaciones.
"""
import os
import re
import math
import time
import subprocess
from typing import List, Tuple, Dict, Any, Optional, Callable
from config import logger
# Expresiones regulares idénticas a internal/ffmpeg/ffmpeg.go
TIME_REGEX = re.compile(r"time=(\d+):(\d+):(\d+\.\d+)")
SPEED_REGEX = re.compile(r"speed=\s*([\d\.]+)x")
PRESET_RES_REGEX = re.compile(r"^\d{2,4}p?$")
CUSTOM_RES_REGEX = re.compile(r"^\d{2,5}[x:]\d{2,5}$")
POS_NUM_REGEX = re.compile(r"^\d{1,3}(\.\d{1,2})?$")
BITRATE_REGEX = re.compile(r"^\d{1,4}k?$")
PAD_FILTER = "pad=ceil(iw/2)*2:ceil(ih/2)*2"
VALID_CODECS = {
"libx264": True, "libx265": True, "h264_qsv": True, "hevc_qsv": True,
"h264_nvenc": True, "hevc_nvenc": True, "h264_amf": True, "hevc_amf": True,
"libvpx-vp9": True, "libaom-av1": True, "libsvtav1": True, "mpeg4": True, "copy": True,
}
VALID_PRESETS = {
"ultrafast": True, "superfast": True, "veryfast": True, "faster": True,
"fast": True, "medium": True, "slow": True, "slower": True, "veryslow": True,
}
SUPPORTED_SUBTITLE_CODECS = {
"subrip": True, "srt": True, "ass": True, "ssa": True,
"mov_text": True, "webvtt": True, "text": True,
}
def get_duration(file_path: str, timeout: float = 30.0) -> float:
"""Obtiene la duración en segundos del video usando ffprobe con timeout de 30s."""
try:
cmd = [
"ffprobe",
"-v", "error",
"-show_entries", "format=duration",
"-of", "default=noprint_wrappers=1:nokey=1",
file_path,
]
res = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
if res.returncode != 0:
logger.warning(f"ffprobe retorno no-cero ({res.returncode}): {res.stderr}")
return 0.0
val_str = res.stdout.strip()
return float(val_str)
except Exception as e:
logger.warning(f"Error ejecutando ffprobe para duración: {e}")
return 0.0
def find_font_file() -> str:
"""Busca una fuente TTF válida en el sistema para drawtext."""
candidates = [
"/usr/share/fonts/dejavu/DejaVuSans-Bold.ttf",
"/usr/share/fonts/dejavu/DejaVuSans.ttf",
"/usr/share/fonts/TTF/DejaVuSans.ttf",
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
"/usr/share/fonts/ttf-dejavu/DejaVuSans.ttf",
"/usr/share/fonts/noto/NotoSans-Regular.ttf",
"C:\\Windows\\Fonts\\arial.ttf",
]
for c in candidates:
if os.path.isfile(c):
return c
return ""
def escape_font_path(p: str) -> str:
"""Adapta una ruta de fuente al filtergraph: forward slashes y sin comillas."""
return p.replace("\\", "/").replace("'", "")
def get_text_subtitle_streams(file_path: str, timeout: float = 30.0) -> List[int]:
"""Detecta las pistas de subtítulos de texto soportadas (srt, ass, ssa, mov_text, webvtt)."""
try:
cmd = [
"ffprobe",
"-v", "error",
"-select_streams", "s",
"-show_entries", "stream=index,codec_name",
"-of", "csv=p=0",
file_path,
]
res = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
if res.returncode != 0:
return []
stream_indices = []
for line in res.stdout.strip().splitlines():
line = line.strip()
if not line:
continue
parts = line.split(",")
if len(parts) >= 2:
try:
idx = int(parts[0].strip())
codec = parts[1].strip().lower()
if SUPPORTED_SUBTITLE_CODECS.get(codec, False):
stream_indices.append(idx)
except ValueError:
continue
return stream_indices
except Exception:
return []
def is_valid_resolution(res: str) -> bool:
"""Indica si la cadena tiene un formato de resolución aceptado."""
res = (res or "").strip()
if not res or res.lower() == "original":
return True
lower = res.lower()
return bool(PRESET_RES_REGEX.match(lower) or CUSTOM_RES_REGEX.match(lower))
def build_scale_filter(res: str) -> str:
"""
Genera el filtro de video adaptado inteligentemente a videos horizontales y verticales.
Idéntico a BuildScaleFilter en internal/ffmpeg/ffmpeg.go.
"""
res = (res or "").strip()
if not res or res.lower() == "original":
return PAD_FILTER
lower = res.lower()
if PRESET_RES_REGEX.match(lower):
clean_num = lower.rstrip("p")
try:
h = int(clean_num)
if 0 < h <= 4320:
# Escala inteligente según orientación:
# Horizontal (iw >= ih): ancho automático -2, altura H
# Vertical / Shorts / Reels (iw < ih): ancho H, altura automática -2
return f"scale='if(gte(iw,ih),-2,{h})':'if(gte(iw,ih),{h},-2)',{PAD_FILTER}"
except ValueError:
pass
return PAD_FILTER
if CUSTOM_RES_REGEX.match(lower):
clean_res = lower.replace("x", ":")
return f"scale={clean_res}:force_original_aspect_ratio=decrease,{PAD_FILTER}"
return PAD_FILTER
def sanitize_watermark(text: str) -> str:
"""
Deja el texto de la marca de agua en un conjunto de caracteres inofensivo
para el filtergraph de FFmpeg (máximo 60 runas).
"""
if not text:
return ""
allowed_chars = set(" -_.()!?¡¿#@&+=")
cleaned = []
for ch in text:
if ch.isalnum() or ch in allowed_chars or ("\u00C0" <= ch <= "\u02AF"):
cleaned.append(ch)
else:
cleaned.append(" ")
res = "".join(cleaned).strip()
return res[:60]
def sanitize_settings(s: Dict[str, Any]) -> Dict[str, Any]:
"""
Sanitiza los ajustes del video contra sus dominios válidos.
Defensa en profundidad idéntica a SanitizeSettings en Go.
"""
sanitized = dict(s)
sanitized["resolution"] = str(s.get("resolution", "")).strip()
sanitized["watermark"] = sanitize_watermark(str(s.get("watermark", "")))
fmt = str(s.get("format", "")).strip().lower()
sanitized["format"] = "avi" if fmt == "avi" else "mp4"
codec = str(s.get("codec", "")).strip().lower()
sanitized["codec"] = codec if VALID_CODECS.get(codec, False) else ""
preset = str(s.get("preset", "")).strip().lower()
sanitized["preset"] = preset if VALID_PRESETS.get(preset, False) else ""
crf = str(s.get("crf", "")).strip()
if crf:
if not POS_NUM_REGEX.match(crf):
sanitized["crf"] = ""
else:
try:
v = float(crf)
sanitized["crf"] = crf if v <= 63 else ""
except ValueError:
sanitized["crf"] = ""
else:
sanitized["crf"] = ""
fps = str(s.get("fps", "")).strip()
if fps:
if not POS_NUM_REGEX.match(fps):
sanitized["fps"] = ""
else:
try:
v = float(fps)
sanitized["fps"] = fps if 1 <= v <= 120 else ""
except ValueError:
sanitized["fps"] = ""
else:
sanitized["fps"] = ""
ab = str(s.get("audio_bitrate", "")).strip().lower()
sanitized["audio_bitrate"] = ab if BITRATE_REGEX.match(ab) else ""
return sanitized
def wm_text_position(pos: str) -> Tuple[str, str]:
"""Posición x/y de drawtext según la posición elegida."""
p = (pos or "tr").strip().lower()
if p == "tl":
return "10", "10"
elif p == "br":
return "w-tw-10", "h-th-10"
elif p == "bl":
return "10", "h-th-10"
return "w-tw-10", "10"
def wm_overlay_expr(pos: str) -> str:
"""Expresión de overlay para marca de agua en imagen."""
p = (pos or "tr").strip().lower()
if p == "tl":
return "10:10"
elif p == "br":
return "W-w-10:H-h-10"
elif p == "bl":
return "10:H-h-10"
return "W-w-10:10"
def generate_thumbnail(video_path: str, thumb_path: str, at_seconds: float = 1.0) -> bool:
"""Extrae un frame del video para usarlo como miniatura."""
try:
ss = f"{at_seconds:.2f}"
cmd = ["ffmpeg", "-y", "-ss", ss, "-i", video_path, "-vframes", "1", "-q:v", "2", thumb_path]
res = subprocess.run(cmd, capture_output=True, timeout=20)
return res.returncode == 0
except Exception:
return False
def normalize_thumbnail_image(src_path: str, dst_path: str) -> bool:
"""Normaliza miniatura a JPEG 320x320 con recorte centrado."""
try:
vf = "scale=320:320:force_original_aspect_ratio=increase:flags=lanczos,crop=320:320"
cmd = ["ffmpeg", "-y", "-i", src_path, "-vf", vf, "-frames:v", "1", "-q:v", "3", dst_path]
res = subprocess.run(cmd, capture_output=True, timeout=20)
return res.returncode == 0
except Exception:
return False
def is_subtitle_or_mux_error(err_str: str) -> bool:
"""Comprueba si el error de stderr de FFmpeg está relacionado con subtítulos corruptos o fallos de muxing."""
lower = err_str.lower()
indicators = [
"application provided duration",
"mov_text",
"subtitle",
"av_interleaved_write_frame",
"invalid argument",
"error while writing stream",
"pts has no value",
"out of range for mov/mp4",
"invalid dts",
"invalid pts",
"muxing",
]
return any(ind in lower for ind in indicators)
def build_compress_args(
input_path: str,
output_path: str,
settings: Dict[str, Any],
wm_image_path: str,
include_subtitles: bool = True,
) -> Tuple[List[str], bool]:
"""Genera la lista de argumentos para FFmpeg, permitiendo habilitar o deshabilitar subtítulos."""
settings = sanitize_settings(settings)
args = ["-y", "-i", input_path]
use_image_wm = bool(wm_image_path and os.path.isfile(wm_image_path))
if use_image_wm:
args.extend(["-i", wm_image_path])
chain = (
f"[0:v]{build_scale_filter(settings.get('resolution', ''))}[vbase];"
f"[1:v]scale='if(gt(iw,150),150,iw)':-1[wm];"
f"[vbase][wm]overlay={wm_overlay_expr(settings.get('watermark_pos', ''))}[v]"
)
args.extend(["-filter_complex", chain, "-map", "[v]", "-map", "0:a?"])
else:
vf_parts = [build_scale_filter(settings.get("resolution", ""))]
wm_text = settings.get("watermark", "")
if wm_text:
x_pos, y_pos = wm_text_position(settings.get("watermark_pos", ""))
font_path = find_font_file()
if font_path:
esc_path = escape_font_path(font_path)
vf_parts.append(
f"drawtext=fontfile='{esc_path}':text='{wm_text}':x={x_pos}:y={y_pos}:fontsize=20:fontcolor=white:box=1:boxcolor=black@0.4:boxborderw=4"
)
else:
vf_parts.append(
f"drawtext=font='Sans':text='{wm_text}':x={x_pos}:y={y_pos}:fontsize=20:fontcolor=white"
)
args.extend(["-map", "0:v:0", "-map", "0:a?", "-vf", ",".join(vf_parts)])
is_avi = settings.get("format", "").lower() == "avi" or output_path.lower().endswith(".avi")
has_subtitles = False
# Subtítulos de texto para contenedores no AVI
if not is_avi and include_subtitles:
text_sub_streams = get_text_subtitle_streams(input_path)
for idx in text_sub_streams:
args.extend(["-map", f"0:{idx}"])
if text_sub_streams:
args.extend(["-c:s", "mov_text"])
has_subtitles = True
codec = settings.get("codec") or "libx264"
args.extend(["-c:v", codec])
if settings.get("preset") and codec != "copy":
args.extend(["-preset", settings.get("preset")])
if settings.get("crf") and codec != "copy":
args.extend(["-crf", str(settings.get("crf"))])
if settings.get("fps") and codec != "copy":
args.extend(["-r", str(settings.get("fps"))])
audio_bitrate = settings.get("audio_bitrate") or "128k"
clean_bitrate = audio_bitrate.lower().rstrip("k")
ac_channels = "2"
try:
if int(clean_bitrate) <= 64:
ac_channels = "1"
except ValueError:
pass
audio_codec = "libmp3lame" if is_avi else "aac"
args.extend([
"-pix_fmt", "yuv420p",
"-c:a", audio_codec,
"-b:a", audio_bitrate,
"-ac", ac_channels,
])
if not is_avi:
args.extend(["-movflags", "+faststart"])
args.extend(["-threads", "0", output_path])
return args, has_subtitles
def execute_ffmpeg_cmd(
args: List[str],
total_duration: float,
on_cmd: Optional[Callable[[subprocess.Popen], None]] = None,
on_progress: Optional[Callable[[Dict[str, Any]], None]] = None,
is_cancelled: Optional[Callable[[], bool]] = None,
) -> None:
cmd_list = ["ffmpeg"] + args
proc = subprocess.Popen(
cmd_list,
stdout=subprocess.DEVNULL,
stderr=subprocess.PIPE,
bufsize=0,
)
if on_cmd:
on_cmd(proc)
last_report_time = 0.0
last_stderr_lines = []
buffer = bytearray()
# Lector que procesa tanto \r como \n (scanCRLFOrCR idéntico a Go)
try:
while True:
if is_cancelled and is_cancelled():
proc.kill()
raise InterruptedError("compresión cancelada")
chunk = proc.stderr.read(128)
if not chunk:
if buffer:
line = buffer.decode("utf-8", errors="replace").strip()
if line:
last_stderr_lines.append(line)
break
for byte in chunk:
if byte in (10, 13): # \n o \r
if buffer:
line = buffer.decode("utf-8", errors="replace").strip()
buffer.clear()
if line:
last_stderr_lines.append(line)
if len(last_stderr_lines) > 30:
last_stderr_lines.pop(0)
matches = TIME_REGEX.search(line)
if matches:
h = float(matches.group(1))
m = float(matches.group(2))
s = float(matches.group(3))
cur_time_sec = h * 3600 + m * 60 + s
percent = 0.0
if total_duration > 0:
percent = min(100.0, (cur_time_sec / total_duration) * 100.0)
speed = "1.0x"
sp_match = SPEED_REGEX.search(line)
if sp_match:
speed = sp_match.group(1) + "x"
eta_str = ""
if total_duration > 0 and cur_time_sec > 0:
rem_sec = max(0.0, total_duration - cur_time_sec)
try:
sp_val = float(speed.rstrip("x"))
if sp_val > 0:
real_rem = int(rem_sec / sp_val)
eta_str = f"{real_rem // 3600:02d}:{(real_rem % 3600) // 60:02d}:{real_rem % 60:02d}"
except ValueError:
pass
time_str = f"{int(h):02d}:{int(m):02d}:{int(s):02d}"
cur_now = time.time()
if cur_now - last_report_time >= 1.0:
last_report_time = cur_now
if on_progress:
on_progress({
"percent": round(percent, 1),
"speed": speed,
"time": time_str,
"eta": eta_str,
})
else:
buffer.append(byte)
proc.wait()
if proc.returncode != 0:
if is_cancelled and is_cancelled():
raise InterruptedError("compresión cancelada")
detail = "\n".join(last_stderr_lines)
raise RuntimeError(f"ffmpeg falló con código {proc.returncode}:\n{detail}")
finally:
if proc.poll() is None:
try:
proc.kill()
except Exception:
pass
def compress_with_cmd(
input_path: str,
output_path: str,
settings: Dict[str, Any],
wm_image_path: str,
total_duration: float,
on_cmd: Optional[Callable[[subprocess.Popen], None]] = None,
on_progress: Optional[Callable[[Dict[str, Any]], None]] = None,
is_cancelled: Optional[Callable[[], bool]] = None,
) -> None:
"""
Ejecuta la compresión con FFmpeg reportando el progreso en vivo.
Si la compresión falla debido a subtítulos corruptos o timestamps inválidos,
reintenta automáticamente descartando los subtítulos para no fallar la tarea.
"""
args, has_subtitles = build_compress_args(input_path, output_path, settings, wm_image_path, include_subtitles=True)
try:
execute_ffmpeg_cmd(args, total_duration, on_cmd, on_progress, is_cancelled)
return
except RuntimeError as e:
if is_cancelled and is_cancelled():
raise
if has_subtitles and is_subtitle_or_mux_error(str(e)):
logger.warning(f"[ffmpeg] Falló compresión con subtítulos ({e}). Reintentando sin subtítulos...")
if os.path.isfile(output_path):
try:
os.remove(output_path)
except Exception:
pass
retry_args, _ = build_compress_args(input_path, output_path, settings, wm_image_path, include_subtitles=False)
execute_ffmpeg_cmd(retry_args, total_duration, on_cmd, on_progress, is_cancelled)
logger.info("[ffmpeg] Compresión completada con éxito tras descartar subtítulos corruptos.")
return
raise