app.py

import os
import asyncio
import re
import json
import subprocess
import tempfile
import shutil
from pathlib import Path
from typing import List, Optional
from datetime import datetime

import cv2
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse, Response
from paddleocr import PaddleOCR
import imageio_ffmpeg

# 获取内置 ffmpeg 路径
FFMPEG_PATH = imageio_ffmpeg.get_ffmpeg_exe()

app = FastAPI(title="视频智能分割工具")

# 全局变量
ocr_engine = None
TEMP_DIR = Path(tempfile.mkdtemp(prefix="video_splitter_"))
OUTPUT_DIR = TEMP_DIR / "output"
OUTPUT_DIR.mkdir(exist_ok=True)

def formatTime(seconds):
    """格式化时间"""
    m = int(seconds // 60)
    s = int(seconds % 60)
    return f"{m:02d}:{s:02d}"

def get_ocr():
    """延迟初始化OCR引擎"""
    global ocr_engine
    if ocr_engine is None:
        print("正在初始化 PaddleOCR 引擎...")
        ocr_engine = PaddleOCR(
            use_angle_cls=True,
            lang='ch',
            use_gpu=False,
            show_log=False
        )
        print("OCR 引擎初始化完成")
    return ocr_engine

def sanitize_filename(text: str) -> str:
    """清理文件名中的非法字符"""
    illegal_chars = r'[<>:"/\\|?*\x00-\x1f]'
    text = re.sub(illegal_chars, '_', text)
    text = text.strip(' .')
    return text[:50] if text else "未命名"

def detect_title_changes(ocr_results: List[dict], similarity_threshold: float = 0.7):
    """检测标题变化点"""
    from difflib import SequenceMatcher

    def is_digit_changed(text1: str, text2: str) -> bool:
        """检测两个文本是否只是数字部分发生变化"""
        text1 = text1.lower()
        text2 = text2.lower()

        nums1 = re.findall(r'\d+', text1)
        nums2 = re.findall(r'\d+', text2)

        if nums1 and nums2 and nums1 != nums2:
            non_digit1 = re.sub(r'\d+', '', text1)
            non_digit2 = re.sub(r'\d+', '', text2)
            if non_digit1 and non_digit2:
                sim = SequenceMatcher(None, non_digit1, non_digit2).ratio()
                if sim > 0.8:
                    return True
        return False

    def calc_similarity(text1: str, text2: str) -> float:
        if not text1 or not text2:
            return 0

        if is_digit_changed(text1, text2):
            return 0

        return SequenceMatcher(None, text1, text2).ratio()

    segments = []
    if not ocr_results:
        return segments

    current_title = ocr_results[0]['text']
    start_time = ocr_results[0]['timestamp']

    for i in range(1, len(ocr_results)):
        text = ocr_results[i]['text']
        timestamp = ocr_results[i]['timestamp']

        similarity = calc_similarity(current_title, text)

        if similarity < similarity_threshold and text.strip():
            segments.append({
                "title": current_title or "未识别标题",
                "start_time": start_time,
                "end_time": timestamp,
                "change_detected": True,
                "from_title": current_title,
                "to_title": text
            })
            current_title = text
            start_time = timestamp

    if ocr_results:
        segments.append({
            "title": current_title or "未识别标题",
            "start_time": start_time,
            "end_time": ocr_results[-1]['timestamp']
        })

    return segments



@app.post("/api/upload")
async def upload_video(file: UploadFile = File(...)):
    if not file.filename:
        raise HTTPException(status_code=400, detail="未选择文件")

    ext = Path(file.filename).suffix
    video_path = TEMP_DIR / f"input{ext}"

    with open(video_path, "wb") as f:
        content = await file.read()
        f.write(content)

    cap = cv2.VideoCapture(str(video_path))
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps if fps > 0 else 0
    cap.release()

    return {
        "success": True,
        "video_id": "input",
        "filename": file.filename,
        "info": {
            "fps": round(fps, 2),
            "width": width,
            "height": height,
            "duration": round(duration, 2),
            "total_frames": total_frames
        }
    }

def ocr_frame(cap, ocr, x, y, w, h, frame_num, fps):
    """对单帧进行OCR识别"""
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
    ret, frame = cap.read()
    if not ret:
        return ""

    roi_frame = frame[y:y+h, x:x+w]
    try:
        ocr_result = ocr.ocr(roi_frame, cls=True)
        if ocr_result and ocr_result[0]:
            texts = [line[1][0] for line in ocr_result[0]]
            return " ".join(texts).strip()
    except Exception:
        pass
    return ""

@app.post("/api/analyze")
async def analyze_video(
    video_id: str = Form(...),
    roi: str = Form(...),
    sample_interval: float = Form(1.0),
    similarity_threshold: float = Form(0.7)
):
    """分析视频并检测标题变化 - 使用SSE实时返回进度"""
    video_path = None
    for ext in ['.mp4', '.avi', '.mov', '.mkv', '.webm']:
        candidate = TEMP_DIR / f"{video_id}{ext}"
        if candidate.exists():
            video_path = candidate
            break

    if not video_path or not Path(video_path).exists():
        raise HTTPException(status_code=404, detail="视频文件不存在")

    roi_data = json.loads(roi)

    async def event_generator():
        cap = cv2.VideoCapture(str(video_path))
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = total_frames / fps if fps > 0 else 0

        x, y, w, h = roi_data['x'], roi_data['y'], roi_data['width'], roi_data['height']
        ocr = get_ocr()

        # ========== 第一轮：2秒间隔快速扫描 ==========
        first_pass_results = []
        scan_interval = 2.0  # 2秒间隔
        current_time = 0.0

        while current_time < duration:
            frame_num = int(current_time * fps)
            text = ocr_frame(cap, ocr, x, y, w, h, frame_num, fps)

            first_pass_results.append({
                "timestamp": round(current_time, 2),
                "text": text
            })

            log_msg = f"[第一轮] {formatTime(current_time)}: {text if text else '无文字'}"
            data = {"time": formatTime(current_time), "log": log_msg, "progress": current_time / duration * 0.5}
            yield "data: " + json.dumps(data) + "\n\n"
            await asyncio.sleep(0)

            current_time += scan_interval

        # 第一轮检测标题变化点
        rough_segments = detect_title_changes(first_pass_results, similarity_threshold)

        if not rough_segments:
            cap.release()
            final_data = {"progress": 1.0, "ocr_results": first_pass_results, "segments": []}
            yield "data: " + json.dumps(final_data) + "\n\n"
            await asyncio.sleep(0)
            return

        # ========== 第二轮：对每个变化点用二分法精确化 ==========
        refined_results = []
        total_rounds = len(rough_segments)
        search_range = 2.0  # ±2秒范围
        precision = 0.1  # 0.1秒精度

        for idx, seg in enumerate(rough_segments):
            change_time = seg.get("start_time", 0)

            # 二分法搜索精确切分点
            low = max(0, change_time - search_range)
            high = min(duration, change_time + search_range)

            prev_text = ""
            best_time = change_time

            while high - low > precision:
                mid = (low + high) / 2
                mid_frame = int(mid * fps)
                mid_text = ocr_frame(cap, ocr, x, y, w, h, mid_frame, fps)

                # 检查与前一个的变化
                if prev_text:
                    from difflib import SequenceMatcher
                    sim = SequenceMatcher(None, prev_text, mid_text).ratio()
                    if sim < similarity_threshold:
                        high = mid
                    else:
                        low = mid
                else:
                    low = mid

                prev_text = mid_text

            final_time = (low + high) / 2

            # 获取该精确时间的文字
            final_frame = int(final_time * fps)
            final_text = ocr_frame(cap, ocr, x, y, w, h, final_frame, fps)

            refined_results.append({
                "timestamp": round(final_time, 2),
                "text": final_text
            })

            log_msg = f"[第二轮] 精确化片段 {idx+1}: {formatTime(final_time)} - {final_text if final_text else '无文字'}"
            data = {"time": formatTime(final_time), "log": log_msg, "progress": 0.5 + (idx / total_rounds) * 0.4}
            yield "data: " + json.dumps(data) + "\n\n"
            await asyncio.sleep(0)

        # 合并两轮结果用于最终片段检测
        all_results = first_pass_results + refined_results
        all_results.sort(key=lambda x: x["timestamp"])

        segments = detect_title_changes(all_results, similarity_threshold)

        cap.release()

        final_data = {"progress": 1.0, "ocr_results": all_results, "segments": segments}
        yield "data: " + json.dumps(final_data) + "\n\n"
        await asyncio.sleep(0)

    return StreamingResponse(event_generator(), media_type="text/event-stream")

@app.post("/api/split")
async def split_video(
    video_id: str = Form(...),
    segments: str = Form(...),
    output_format: str = Form("mp4")
):
    """切分视频 - 使用SSE返回进度"""
    video_path = None
    for ext in ['.mp4', '.avi', '.mov', '.mkv', '.webm']:
        candidate = TEMP_DIR / f"{video_id}{ext}"
        if candidate.exists():
            video_path = candidate
            break

    if not video_path:
        raise HTTPException(status_code=404, detail="视频文件不存在")

    segments_data = json.loads(segments)

    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    OUTPUT_DIR.mkdir()

    async def event_generator():
        output_files = []
        total = len(segments_data)

        for i, seg in enumerate(segments_data):
            title = sanitize_filename(seg['title'])

            if output_format == "mp3":
                output_name = f"{i+1:02d}_{title}.mp3"
                output_path = os.path.join(str(OUTPUT_DIR), output_name)
            else:
                output_name = f"{i+1:02d}_{title}.mp4"
                output_path = os.path.join(str(OUTPUT_DIR), output_name)

            start = seg['start_time']
            duration = seg['end_time'] - seg['start_time']

            log_msg = ""
            try:
                if output_format == "mp3":
                    cmd = [
                        FFMPEG_PATH, '-y',
                        '-ss', str(start),
                        '-i', str(video_path),
                        '-t', str(duration),
                        '-vn',
                        '-acodec', 'libmp3lame',
                        '-q:a', '2',
                        output_path
                    ]
                else:
                    cmd = [
                        FFMPEG_PATH, '-y',
                        '-ss', str(start),
                        '-i', str(video_path),
                        '-t', str(duration),
                        '-c', 'copy',
                        '-avoid_negative_ts', 'make_zero',
                        output_path
                    ]

                subprocess.run(cmd, capture_output=True, check=True)
                output_files.append({
                    "name": output_name,
                    "path": output_path,
                    "title": seg['title'],
                    "start": seg['start_time'],
                    "end": seg['end_time']
                })
                log_msg = f"已生成: {output_name}"

            except Exception as e:
                log_msg = f"切分失败: {output_name}, 错误: {e}"
                print(log_msg)

            # 发送SSE进度
            progress = (i + 1) / total
            data = {"progress": progress, "log": log_msg, "time": f"{i+1}/{total}"}
            yield "data: " + json.dumps(data) + "\n\n"
            await asyncio.sleep(0)

        # 返回最终结果
        final_data = {"progress": 1.0, "files": output_files, "success": True}
        yield "data: " + json.dumps(final_data) + "\n\n"
        await asyncio.sleep(0)

    return StreamingResponse(event_generator(), media_type="text/event-stream")

@app.get("/api/frame/{video_id}")
async def get_frame(video_id: str, frame: int = 0):
    video_path = None
    for ext in ['.mp4', '.avi', '.mov', '.mkv', '.webm']:
        candidate = TEMP_DIR / f"{video_id}{ext}"
        if candidate.exists():
            video_path = candidate
            break

    if not video_path:
        raise HTTPException(status_code=404, detail="视频文件不存在")

    cap = cv2.VideoCapture(str(video_path))
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame)
    ret, frame_data = cap.read()
    cap.release()

    if not ret:
        raise HTTPException(status_code=400, detail="无法读取帧")

    success, buffer = cv2.imencode('.jpg', frame_data, [cv2.IMWRITE_JPEG_QUALITY, 85])
    if not success:
        raise HTTPException(status_code=500, detail="编码失败")

    return Response(content=buffer.tobytes(), media_type="image/jpeg")

@app.get("/download/{filename}")
async def download_file(filename: str):
    file_path = OUTPUT_DIR / filename
    if not file_path.exists():
        raise HTTPException(status_code=404, detail="文件不存在")
    return FileResponse(file_path, filename=filename)

@app.get("/download_all")
async def download_all():
    import zipfile

    mp4_files = list(OUTPUT_DIR.glob("*.mp4"))
    mp3_files = list(OUTPUT_DIR.glob("*.mp3"))

    if mp3_files:
        zip_name = "split_audios.zip"
        files = mp3_files
    else:
        zip_name = "split_videos.zip"
        files = mp4_files

    zip_path = TEMP_DIR / zip_name

    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
        for file in files:
            zf.write(file, file.name)

    return FileResponse(zip_path, filename=zip_name)

@app.get("/")
async def root():
    return FileResponse("static/index.html")

@app.on_event("startup")
async def startup_event():
    import webbrowser
    import threading
    import time

    def open_browser():
        time.sleep(2)
        webbrowser.open("http://localhost:8000")

    threading.Thread(target=open_browser, daemon=True).start()

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)