# -*- coding: utf-8 -*-
|
import json
|
import logging
|
import os
|
import shutil
|
import subprocess
|
import tempfile
|
from contextlib import contextmanager
|
from typing import Optional
|
|
import httpx
|
|
logger = logging.getLogger(__name__)
|
|
|
def get_ffmpeg_cmd(name: str = "ffmpeg") -> str:
|
ffmpeg_dir = "D:/code/ffmpeg/"
|
# ffmpeg_dir = os.environ.get("FFMPEG_DIR", "")
|
print(f"ffmpeg_dir:{ffmpeg_dir}")
|
if ffmpeg_dir:
|
exe = "ffmpeg.exe" if os.name == "nt" else "ffmpeg"
|
print(f"exe:{exe}")
|
if name == "ffprobe":
|
exe = "ffprobe.exe" if os.name == "nt" else "ffprobe"
|
return os.path.join(ffmpeg_dir, exe)
|
return name
|
|
|
def probe_duration(video_path: str) -> float:
|
cmd = [
|
get_ffmpeg_cmd("ffprobe"),
|
"-v", "error",
|
"-show_entries", "format=duration",
|
"-of", "default=noprint_wrappers=1:nokey=1",
|
video_path,
|
]
|
try:
|
out = subprocess.check_output(cmd, stderr=subprocess.DEVNULL, text=True).strip()
|
return float(out) if out else 0.0
|
except Exception as e:
|
logger.warning("ffprobe 失败: %s", e)
|
return 0.0
|
|
|
def download_video(url: str, dest_path: str, timeout: float = 600.0) -> None:
|
with httpx.stream("GET", url, timeout=timeout, follow_redirects=True) as resp:
|
resp.raise_for_status()
|
with open(dest_path, "wb") as f:
|
for chunk in resp.iter_bytes(chunk_size=65536):
|
f.write(chunk)
|
|
|
@contextmanager
|
def temp_video(video_url: Optional[str], duration_hint: float = 0.0):
|
tmp_dir = tempfile.mkdtemp(prefix="snap_infer_")
|
video_path = os.path.join(tmp_dir, "video.mp4")
|
try:
|
if video_url:
|
download_video(video_url, video_path)
|
else:
|
raise ValueError("video_url 不能为空")
|
duration = probe_duration(video_path)
|
if duration <= 0:
|
duration = duration_hint if duration_hint > 0 else 1200.0
|
yield video_path, duration
|
finally:
|
shutil.rmtree(tmp_dir, ignore_errors=True)
|