app.py

# 新增依赖

import gradio as gr
import cv2
import numpy as np
import os
from datetime import datetime  # 新增时间模块

from scenedetect import open_video, SceneManager
from scenedetect.detectors import ContentDetector
from moviepy.editor import VideoFileClip

import random
from functools import partial

import clip
import decord
import nncore
import torch
import torchvision.transforms.functional as F
from decord import VideoReader
from nncore.engine import load_checkpoint
from nncore.nn import build_model

import pandas as pd

def convert_time(seconds):
    minutes, seconds = divmod(round(max(seconds, 0)), 60)
    return f'{minutes:02d}:{seconds:02d}'

# R2-Tuning 模型配置
TUNING_CONFIG = 'configs/qvhighlights/r2_tuning_qvhighlights.py'
TUNING_WEIGHT = 'https://huggingface.co/yeliudev/R2-Tuning/resolve/main/checkpoints/r2_tuning_qvhighlights-ed516355.pth'

# 初始化R2-Tuning模型
def init_tuning_model(config, checkpoint):
    cfg = nncore.Config.from_file(config)
    cfg.model.init = True

    if checkpoint.startswith('http'):
        checkpoint = nncore.download(checkpoint, out_dir='checkpoints')

    model = build_model(cfg.model, dist=False).eval()
    model = load_checkpoint(model, checkpoint, warning=False)
    return model, cfg

tuning_model, tuning_cfg = init_tuning_model(TUNING_CONFIG, TUNING_WEIGHT)

# 视频预处理函数（来自第一个应用）
def preprocess_video(video_path, cfg):
    decord.bridge.set_bridge('torch')
    vr = decord.VideoReader(video_path)
    stride = vr.get_avg_fps() / cfg.data.val.fps
    fm_idx = [min(round(i), len(vr) - 1) for i in np.arange(0, len(vr), stride).tolist()]
    video = vr.get_batch(fm_idx).permute(0, 3, 1, 2).float() / 255

    size = 336 if '336px' in cfg.model.arch else 224
    h, w = video.size(-2), video.size(-1)
    s = min(h, w)
    x, y = round((h - s) / 2), round((w - s) / 2)
    video = video[..., x:x + s, y:y + s]
    video = F.resize(video, size=(size, size))
    video = F.normalize(video, (0.481, 0.459, 0.408), (0.269, 0.261, 0.276))
    return video.reshape(video.size(0), -1).unsqueeze(0)

# 在calculate_saliency函数中增加视频时长返回
def calculate_saliency(video_path, query, model, cfg):
    if len(query) == 0:
        return None, None, 0
    
    video = preprocess_video(video_path, cfg)
    query = clip.tokenize(query, truncate=True)
    
    device = next(model.parameters()).device
    data = dict(video=video.to(device), query=query.to(device), fps=[cfg.data.val.fps])
    
    with torch.inference_mode():
        pred = model(data)
    
    hd = pred['_out']['saliency'].cpu()
    hd = ((hd - hd.min()) / (hd.max() - hd.min()) * 0.9 + 0.05).numpy()
    time_axis = np.arange(0, len(hd) * 2, 2)
    
    # 获取视频总时长
    vr = decord.VideoReader(video_path)
    duration = len(vr) / vr.get_avg_fps()
    return hd, time_axis, duration

# 修改后的场景检测函数
def find_scenes(video_path, threshold, query):
    # 初始化输出目录
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    output_dir = f"output_{timestamp}"
    os.makedirs(output_dir, exist_ok=True)

    # 计算saliency score
    saliency_scores, time_points, total_duration = calculate_saliency(video_path, query, tuning_model, tuning_cfg)
    if saliency_scores is None:
        raise gr.Error("请输入有效的文本查询")

    # 线性插值处理
    new_time = np.linspace(0, total_duration, num=int(total_duration*10))  # 每0.5秒一个点
    interp_scores = np.interp(new_time, time_points, saliency_scores)

    # 场景检测（原逻辑）
    filename = os.path.splitext(os.path.basename(video_path))[0]
    video = open_video(video_path)
    scene_manager = SceneManager()
    scene_manager.add_detector(ContentDetector(threshold=threshold))
    scene_manager.detect_scenes(video, show_progress=True)
    scene_list = scene_manager.get_scene_list()

    if not scene_list:
        gr.Warning("No scenes detected in this video")
        return None, None, None, None

    # 处理每个场景片段
    processed_scenes = []
    for i, shot in enumerate(scene_list):
        # 获取时间信息
        start_sec = shot[0].get_seconds()
        end_sec = shot[1].get_seconds()
        
        # 使用插值后的得分计算
        start_idx = np.searchsorted(new_time, start_sec, side='left')
        end_idx = np.searchsorted(new_time, end_sec, side='right')
        valid_scores = interp_scores[start_idx:end_idx]
        
        # 处理可能的空值（理论上插值后不会出现）
        valid_scores = valid_scores[~np.isnan(valid_scores)]
        scene_score = valid_scores.mean() if len(valid_scores) > 0 else 0.0
        
        # 保存片段信息
        scene_info = {
            "start": convert_time(start_sec),
            "end": convert_time(end_sec),
            "score": round(float(scene_score), 3),
            "start_sec": start_sec,
            "end_sec": end_sec
        }
        processed_scenes.append(scene_info)

    # 按得分排序
    processed_scenes.sort(key=lambda x: x['score'], reverse=True)

    # 生成输出内容
    timecodes = [{"title": filename + ".mp4", "fps": scene_list[0][0].get_framerate()}]
    shots = []
    stills = []
    
    for idx, scene in enumerate(processed_scenes):
        # 生成片段名称
        shot_name = f"shot_{idx+1}_{filename}"
        target_name = os.path.join(output_dir, f"{shot_name}.mp4")
        
        # 分割视频片段
        with VideoFileClip(video_path) as clip:
            subclip = clip.subclip(scene['start_sec'], scene['end_sec'])
            subclip.write_videofile(target_name, 
                                  codec="libx264",
                                  audio_codec="aac",
                                  threads=4,
                                  preset="fast",
                                  ffmpeg_params=["-crf", "23"])

        # 截取缩略图
        vid = cv2.VideoCapture(video_path)
        vid.set(cv2.CAP_PROP_POS_MSEC, scene['start_sec']*1000)
        ret, frame = vid.read()
        img_path = os.path.join(output_dir, f"{shot_name}_screenshot.png")
        cv2.imwrite(img_path, frame)
        vid.release()

        # 保存结果（带片段名称的标签）
        timecodes.append({
            "tc_in": scene['start'],
            "tc_out": scene['end'],
            "score": scene['score'],
            "shot_name": shot_name
        })
        shots.append(target_name)
        stills.append((img_path, f'{shot_name}\nScore: {scene["score"]:.3f}'))

    # 生成折线图数据（使用插值后的数据）
    plot_data = pd.DataFrame({
        'x': new_time,
        'y': interp_scores
    })

    return timecodes, shots, stills, plot_data

# 修改后的界面
with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("""
            # 增强版场景编辑检测
            新增功能：
            1. 输入文本查询分析视频内容相关性
            2. 显示相关性时序折线图
            3. 按相关性得分排序输出片段
            """)
                
        with gr.Row():
            with gr.Column():
                video_input = gr.Video(sources="upload", format="mp4", label="视频输入")
                query_input = gr.Textbox(label="文本查询", placeholder="输入描述视频内容的文本（5-15单词为佳）")
                threshold = gr.Slider(label="场景切换检测阈值", minimum=15.0, maximum=40.0, value=27.0)
                with gr.Row():
                    clear_button = gr.Button("清除")
                    run_button = gr.Button("开始处理", variant="primary")
                plot_output = gr.LinePlot(x='x', y='y', x_title='时间（秒）', 
                                        y_title='相关性得分', label='时序相关性分析')
            with gr.Column():
                json_output = gr.JSON(label="场景分析结果（按得分排序）")

                file_output = gr.File(label="分割片段下载")
                gallery_output = gr.Gallery(label="场景缩略图", object_fit="cover", columns=3)

        run_button.click(
            fn=find_scenes,
            inputs=[video_input, threshold, query_input],
            outputs=[json_output, file_output, gallery_output, plot_output]
        )
        clear_button.click(
            fn=lambda: [None, 27, None, None, None, None],
            inputs=None,
            outputs=[video_input, threshold, query_input, json_output, file_output, gallery_output]
        )

    gr.Examples(
            examples=[
                ["anime_kiss.mp4", 27, "A romantic kiss scene between two characters"],
                ["anime_tear.mp4", 30, "An anime character is crying."]     
                     ],
            inputs=[video_input, threshold, query_input],
            outputs=[json_output, file_output, gallery_output, plot_output],
            fn=find_scenes,
            cache_examples=False
        )

demo.queue().launch(debug=True, share=True)