157 lines
5.7 KiB
Python
157 lines
5.7 KiB
Python
"""后台任务路由：视频分析（ffmpeg 抽帧 + 视觉模型）。"""
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import re
|
|
import subprocess
|
|
import uuid
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import imageio_ffmpeg
|
|
from fastapi import APIRouter, Depends, HTTPException
|
|
from pydantic import BaseModel
|
|
from sqlalchemy.orm import Session
|
|
|
|
from ..config import settings
|
|
from ..core import llm
|
|
from ..core.deps import get_current_user
|
|
from ..core.response import ok
|
|
from ..database import get_db
|
|
from ..models import FileEntry, Task, User
|
|
|
|
router = APIRouter(prefix="/tasks", tags=["后台任务"])
|
|
|
|
FFMPEG = imageio_ffmpeg.get_ffmpeg_exe()
|
|
_background_tasks: dict[int, asyncio.Task] = {}
|
|
|
|
|
|
def _task_out(t: Task) -> dict:
|
|
try:
|
|
result = json.loads(t.result) if t.result else {}
|
|
params = json.loads(t.params) if t.params else {}
|
|
except json.JSONDecodeError:
|
|
result, params = {}, {}
|
|
return {"id": t.id, "type": t.type, "status": t.status, "params": params, "result": result,
|
|
"error": t.error, "created_at": t.created_at.isoformat(),
|
|
"finished_at": t.finished_at.isoformat() if t.finished_at else None}
|
|
|
|
|
|
def _get_owned(db: Session, user: User, task_id: int) -> Task:
|
|
t = db.get(Task, task_id)
|
|
if not t or t.user_id != user.id:
|
|
raise HTTPException(status_code=404, detail="任务不存在")
|
|
return t
|
|
|
|
|
|
def _db_refresh(db: Session):
|
|
"""后台协程中更新任务状态的辅助。"""
|
|
db.commit()
|
|
|
|
|
|
class VideoAnalysisRequest(BaseModel):
|
|
file_id: int
|
|
prompt: str = "请全面分析这段视频：画面内容、构图、运镜、节奏、亮点与改进建议。"
|
|
frames: int = 8
|
|
|
|
|
|
def _extract_frames(video_path: Path, work_dir: Path, frames: int) -> list[Path]:
|
|
"""用 ffmpeg 均匀抽帧，返回帧文件列表。"""
|
|
work_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 解析视频时长（ffprobe 不可用，用 ffmpeg -i 的 stderr 输出）
|
|
duration = 10.0
|
|
probe = subprocess.run([FFMPEG, "-i", str(video_path)], capture_output=True, text=True)
|
|
m = re.search(r"Duration:\s*(\d+):(\d+):(\d+\.?\d*)", probe.stderr)
|
|
if m:
|
|
duration = int(m.group(1)) * 3600 + int(m.group(2)) * 60 + float(m.group(3))
|
|
|
|
interval = max(1.0, duration / max(frames, 1))
|
|
out_paths: list[Path] = []
|
|
for i in range(frames):
|
|
t = i * interval
|
|
out = work_dir / f"frame_{i:03d}.jpg"
|
|
r = subprocess.run(
|
|
[FFMPEG, "-y", "-ss", str(t), "-i", str(video_path), "-frames:v", "1",
|
|
"-vf", "scale=640:-2", "-q:v", "5", str(out)],
|
|
capture_output=True, text=True,
|
|
)
|
|
if r.returncode == 0 and out.exists():
|
|
out_paths.append(out)
|
|
return out_paths
|
|
|
|
|
|
async def _run_video_analysis(task_id: int, file_id: int, prompt: str, frames: int):
|
|
"""后台执行视频分析（独立 DB 会话，不依赖请求生命周期）。"""
|
|
from ..database import SessionLocal as _SL
|
|
db = _SL()
|
|
try:
|
|
task = db.get(Task, task_id)
|
|
if not task:
|
|
return
|
|
task.status = "running"
|
|
db.commit()
|
|
try:
|
|
f = db.get(FileEntry, file_id)
|
|
video_path = Path(f.stored_path) if f else None
|
|
if not video_path or not video_path.exists():
|
|
raise RuntimeError("视频文件不存在")
|
|
|
|
work_dir = settings.DATA_DIR / "frames" / str(task_id)
|
|
frame_paths = await asyncio.to_thread(_extract_frames, video_path, work_dir, frames)
|
|
if not frame_paths:
|
|
raise RuntimeError("抽帧失败：ffmpeg 可能未安装或视频无法解析")
|
|
|
|
# 帧 → base64 data URL（视觉模型可直接读取）
|
|
image_urls = []
|
|
for fp in frame_paths:
|
|
b64 = base64.b64encode(fp.read_bytes()).decode()
|
|
image_urls.append(f"data:image/jpeg;base64,{b64}")
|
|
|
|
analysis = await llm.vision_analysis(prompt, image_urls)
|
|
task.result = json.dumps({"analysis": analysis, "frame_count": len(frame_paths)},
|
|
ensure_ascii=False)
|
|
task.status = "done"
|
|
except Exception as e:
|
|
task.status = "failed"
|
|
task.error = str(e)
|
|
finally:
|
|
task.finished_at = datetime.utcnow()
|
|
db.commit()
|
|
_background_tasks.pop(task_id, None)
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
@router.get("")
|
|
def list_tasks(user: User = Depends(get_current_user), db: Session = Depends(get_db)):
|
|
tasks = db.query(Task).filter(Task.user_id == user.id).order_by(Task.id.desc()).limit(50).all()
|
|
return ok([_task_out(t) for t in tasks])
|
|
|
|
|
|
@router.get("/{task_id}")
|
|
def get_task(task_id: int, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
|
|
t = _get_owned(db, user, task_id)
|
|
return ok(_task_out(t))
|
|
|
|
|
|
@router.post("/video-analysis")
|
|
async def create_video_analysis(body: VideoAnalysisRequest,
|
|
user: User = Depends(get_current_user), db: Session = Depends(get_db)):
|
|
f = db.get(FileEntry, body.file_id)
|
|
if not f or f.user_id != user.id:
|
|
raise HTTPException(status_code=404, detail="文件不存在")
|
|
if not f.mime.startswith("video"):
|
|
raise HTTPException(status_code=400, detail="请上传视频文件")
|
|
|
|
task = Task(user_id=user.id, type="video_analysis",
|
|
params=json.dumps({"file_id": body.file_id, "prompt": body.prompt, "frames": body.frames},
|
|
ensure_ascii=False))
|
|
db.add(task)
|
|
db.commit()
|
|
db.refresh(task)
|
|
|
|
bg = asyncio.create_task(_run_video_analysis(task.id, body.file_id, body.prompt, body.frames))
|
|
_background_tasks[task.id] = bg
|
|
return ok(_task_out(task))
|