pyai/route/audio.py

import logging
from pathlib import Path

from fastapi import APIRouter, HTTPException
import os
import uuid
import shutil
from fastapi import UploadFile, File
from starlette.concurrency import run_in_threadpool
import service.ai_task as ai_task
import service.pyav as pyav
from setting import UPLOAD_DIR, OUTPUT_DIR

logger = logging.getLogger(__name__)

# 创建路由对象，可以统一设置前缀 (prefix) 和 标签 (tags)
router = APIRouter(
    prefix="/api1/audio",
    tags=["audio"]
)

@router.post("/asr")
async def upload_audio(file: UploadFile = File(...)):
    # 1. 统一生成一次 task_id，确保前后一致
    task_id = str(uuid.uuid4())[:8]
    ext = file.filename.split('.')[-1]
    save_path = os.path.join(UPLOAD_DIR, f"{task_id}.{ext}")

    # 2. 解决 IO 阻塞方案 A: 使用 run_in_threadpool (推荐)
    # 这样会将同步的写入操作丢进单独的线程，不阻塞主事件循环
    def save_file():
        with open(save_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
    await run_in_threadpool(save_file)

    # 3. 构造路径并存入队列
    srt_path = os.path.join(OUTPUT_DIR, f"{task_id}.srt")
    video_path = os.path.join(OUTPUT_DIR, f"{task_id}.mp4")

    # 传递已经确定好的 task_id
    await ai_task.put_task(task_id, save_path, srt_path, video_path)
    return {
        "status": "queued",
        "task_id": task_id,
        "message": "文件已上传并加入 GPU 处理队列",
        "srt_preview_path": f"{OUTPUT_DIR}/{task_id}.srt"
    }

@router.get("/tasks")
async def get_queue_status():
    return {"queue_size": ai_task.get_tasks()}

@router.get("/result/{task_id}")
async def get_asr_result(task_id: str):
    file_name = check_file_prefix(UPLOAD_DIR, task_id)
    if not file_name:
        raise HTTPException(status_code=404, detail="音频文件不存在")
    audio_path = f"{UPLOAD_DIR}/{file_name}"

    txt_path = f"{OUTPUT_DIR}/{task_id}.txt"
    if not os.path.exists(txt_path):
        raise HTTPException(status_code=404, detail="音频文本文件不存在")

    srt_path = f"{OUTPUT_DIR}/{task_id}.srt"
    if not os.path.exists(srt_path):
        raise HTTPException(status_code=404, detail="字幕文件不存在")

    video_path = f"{OUTPUT_DIR}/{task_id}.mp4"
    if not os.path.exists(video_path):
        raise HTTPException(status_code=404, detail="视频文件不存在")

    with open(txt_path, "r", encoding="utf-8") as f:
        text = f.read()

    info = pyav.get_media_info(audio_path)
    srt = pyav.parse_srt_to_list(srt_path)
    return {
        "task_id": task_id,
        "duration": info['duration'],
        "text": text,
        "srt": srt,
        "audio_url": f"/api1/file/audio/{file_name}",
        "video_url": f"/api1/file/video/{task_id}.mp4"
    }


def check_file_prefix(directory, prefix):
    # 1. 转化为 Path 对象
    path = Path(directory)
    # 2. 匹配所有以 prefix 开头的文件
    # 如果要匹配特定后缀，可以使用 f"{prefix}*.jpg"
    matched_files = list(path.glob(f"{prefix}*"))

    count = len(matched_files)
    if count == 1:
        file_path = matched_files[0]
        return file_path.name
    else:
        return None