app.py

import streamlit as st
import os
import time
from PIL import Image

# 全局缓存 Whisper 语音大模型
@st.cache_resource
def load_whisper_model():
    import whisper
    return whisper.load_model("small")

# 页面配置
st.set_page_config(
    page_title="多模态AI生成痕迹鉴别系统",
    page_icon="🔍",
    layout="wide"
)

# ==========================================
# 侧边栏：系统设置与说明
# ==========================================
with st.sidebar:
    st.title("⚙️ 鉴别引擎设置")
    st.markdown("提供轻量、可解释的多模态AI生成内容快速鉴别能力。")
    st.divider()

    st.header("🎚️ 动态判定阈值")
    st.markdown("调整敏感度，适应不同审核场景：")
    high_risk_threshold = st.slider(
        "高危报警阈值",
        min_value=0.60, max_value=0.95, value=0.80, step=0.05,
        help="高于此值判定为“高度疑似AI生成”"
    )
    warning_threshold = st.slider(
        "存疑缓冲阈值",
        min_value=0.20, max_value=0.55, value=0.40, step=0.05,
        help="介于警告阈值与报警阈值之间时，标记为“存疑”"
    )

    st.divider()
    st.caption("🔹 引擎状态：轻量模型已缓存 · 本地推理")

# ==========================================
# 主界面标题
# ==========================================
st.title("🔍 多模态AI生成痕迹快速鉴别系统")
st.markdown(
    "同时支持**图像、文本、视频**单模态扫描，以及**图文/视文**联合跨模态融合鉴定。"
)

work_mode = st.radio(
    "请选择检测模式：",
    ["单模态独立检测（图片 / 文本 / 视频）", "多模态联合检测（图片+文案 / 视频+文案）"],
    horizontal=True
)

st.divider()

# ==========================================
# 模式一：单模态独立检测
# ==========================================
if "单模态独立检测" in work_mode:
    with st.container(border=True):
        uploaded_file = st.file_uploader(
            "拖拽或点击上传待检测文件",
            type=["jpg", "png", "jpeg", "txt", "docx", "mp4", "avi", "mov"],
            help="支持常见图片、文本、视频格式，单文件≤50MB，视频时长建议≤5分钟"
        )

    if uploaded_file is not None:
        file_type = uploaded_file.name.split('.')[-1].lower()
        st.success(f"✅ 已接收：**{uploaded_file.name}** | 启动鉴别引擎...")

        col_content, col_result = st.columns([1, 1.2], gap="large")

        # ----------------- 图像检测 -----------------
        if file_type in ['jpg', 'png', 'jpeg']:
            with col_content:
                st.subheader("原始图像")
                st.image(uploaded_file, use_container_width=True)

                st.subheader("AI痕迹热力图（Grad-CAM）")
                heatmap_placeholder = st.empty()

            with col_result:
                st.subheader("⚙️ 特征提取中...")
                progress_bar = st.progress(0)
                status_text = st.empty()

                for percent_complete in range(100):
                    time.sleep(0.01)
                    progress_bar.progress(percent_complete + 1)
                    status_text.text(f"正在提取LBP纹理、频域伪影及深度语义... {percent_complete + 1}%")

                from image_module import analyze_image, load_deep_image_model, generate_image_heatmap

                result = analyze_image(uploaded_file)

                model, device = load_deep_image_model()
                # 图像模块内部我们已经处理过 seek(0)，所以这里可以直接传
                heatmap_img = generate_image_heatmap(uploaded_file, model, device)
                heatmap_placeholder.image(
                    heatmap_img,
                    caption="🔴 红色区域为模型判定的高可疑伪造区域",
                    use_container_width=True
                )

                status_text.text("特征提取完成")

                st.subheader("检测报告")
                final_prob = result['final_probability']

                if final_prob >= high_risk_threshold:
                    st.error(f"高度疑似AI生成图像（生成概率：{final_prob * 100:.1f}%）")
                elif warning_threshold <= final_prob < high_risk_threshold:
                    st.warning(f"可疑图像，可能经过AI修图或局部生成（概率：{final_prob * 100:.1f}%）")
                else:
                    st.success(f"真实图像，未检出明显生成痕迹（概率：{final_prob * 100:.1f}%）")

                st.progress(float(final_prob))
                with st.expander("展开底层特征参数", expanded=True):
                    st.info(
                        f"**双轨特征得分：**\n\n"
                        f"- 传统物理特征异常度：{result['traditional_score'] * 100:.1f}%\n"
                        f"- MobileNetV2 深度特征：{result['deep_score'] * 100:.1f}%"
                    )

        # ----------------- 文本检测 -----------------
        elif file_type in ['txt', 'docx']:
            text_content = ""
            uploaded_file.seek(0) # 核心修复：倒带文件指针
            if file_type == 'txt':
                text_content = uploaded_file.read().decode("utf-8")
            elif file_type == 'docx':
                import docx
                doc = docx.Document(uploaded_file)
                text_content = "\n".join([para.text for para in doc.paragraphs])

            with col_content:
                st.subheader("原始文本内容")
                text_placeholder = st.empty()
                text_placeholder.text_area("提取的文本", text_content, height=350)

            with col_result:
                st.subheader("⚙️ 语义连贯性分析中...")
                with st.spinner("正在计算困惑度、句法复杂度及BERT深度特征..."):
                    from text_module import analyze_text, get_custom_text_model, generate_text_highlight_html

                    result = analyze_text(text_content)

                    tokenizer, model = get_custom_text_model()
                    highlighted_html = generate_text_highlight_html(text_content, tokenizer, model)
                    text_placeholder.markdown("**🔥 AI痕迹逐句高亮（黄色为高风险句式）**", unsafe_allow_html=True)
                    text_placeholder.markdown(highlighted_html, unsafe_allow_html=True)

                st.subheader("检测报告")
                final_prob = result['final_probability']

                if final_prob >= high_risk_threshold:
                    st.error(f"高度疑似大语言模型生成文本（生成概率：{final_prob * 100:.1f}%）")
                elif warning_threshold <= final_prob < high_risk_threshold:
                    st.warning(f"可疑文本，可能经AI润色或拼接（概率：{final_prob * 100:.1f}%）")
                else:
                    st.success(f"真实人类写作风格，低风险（概率：{final_prob * 100:.1f}%）")

                st.progress(float(final_prob))
                with st.expander("展开底层特征参数", expanded=True):
                    st.info(
                        f"**多维度文本特征：**\n\n"
                        f"- 统计学风格异常度：{result['stat_score'] * 100:.1f}%\n"
                        f"- DistilBERT 语义鉴别得分：{result['deep_score'] * 100:.1f}%\n"
                        f"- 补充指标：{result['details']}"
                    )

        # ----------------- 视频检测 -----------------
        elif file_type in ['mp4', 'avi', 'mov']:
            with col_content:
                st.subheader("视频内容预览")
                st.video(uploaded_file)

            with col_result:
                st.subheader("⚙️ 时空特征提取中...")
                with st.spinner("正在进行关键帧抽帧、光流计算及音频频谱分析..."):
                    temp_video_path = f"temp_uploaded_video.{file_type}"
                    with open(temp_video_path, "wb") as f:
                        uploaded_file.seek(0) # 核心修复：倒带文件指针
                        f.write(uploaded_file.read())

                    from video_module import analyze_video
                    result = analyze_video(temp_video_path)

                    if os.path.exists(temp_video_path):
                        os.remove(temp_video_path)

                if "error" in result:
                    st.error(result["error"])
                else:
                    st.subheader("检测报告")
                    final_prob = result['avg_probability']

                    if final_prob >= high_risk_threshold:
                        st.error(f"整体视频高度疑似AI生成（综合概率：{final_prob * 100:.1f}%）")
                    elif warning_threshold <= final_prob < high_risk_threshold:
                        st.warning(f"视频存疑，存在异常帧或拼接痕迹（概率：{final_prob * 100:.1f}%）")
                    else:
                        st.success(f"未发现明显时序伪造痕迹，低风险（概率：{final_prob * 100:.1f}%）")

                    st.progress(float(final_prob))
                    with st.expander("展开抽帧分析详情", expanded=True):
                        st.info(
                            f"**视频物理特征：**\n\n"
                            f"- 总帧数：{result['total_frames']} 帧 (帧率 {result['fps']:.1f} fps)\n"
                            f"- 关键帧采样数：{result['sampled_frames']} 帧\n"
                            f"- 单帧最高风险值：{result['max_probability'] * 100:.1f}%"
                        )

# ==========================================
# 模式二：多模态联合检测
# ==========================================
elif "多模态联合检测" in work_mode:
    st.subheader("🔗 跨模态联合检测场景")
    st.markdown(
        "模拟真实审核场景：同时检测**图像/视频**与**配套文本**，通过决策级加权融合输出综合AI生成概率。"
    )

    # 初始化session_state记忆体（用于自动提取的文本）
    if 'auto_text' not in st.session_state:
        st.session_state['auto_text'] = ""

    col_media, col_text = st.columns(2, gap="large")

    with col_media:
        multi_media = st.file_uploader(
            "🖼️ / 🎞️ 第一步：上传媒体内容（图片或短视频）",
            type=["jpg", "png", "jpeg", "mp4", "avi", "mov"],
            key="multi_media"
        )
        file_type = ""
        if multi_media:
            file_type = multi_media.name.split('.')[-1].lower()
            if file_type in ['mp4', 'avi', 'mov']:
                st.video(multi_media)

                # 智能语音提取按钮
                if st.button("🎙️ 自动从视频提取语音转文字", use_container_width=True):
                    with st.spinner("正在加载Whisper模型并转录音频（若视频较长请耐心等待数分钟）..."):
                        try:
                            from moviepy import VideoFileClip
                            import whisper

                            # 1. 保存临时视频
                            temp_vid = "temp_for_audio.mp4"
                            with open(temp_vid, "wb") as f:
                                multi_media.seek(0)
                                f.write(multi_media.read())

                            # 2. 读取并剥离音频
                            temp_audio = "temp_audio.wav"
                            my_clip = VideoFileClip(temp_vid)

                            # 【新增防护】检查视频到底有没有声音！
                            if my_clip.audio is None:
                                st.error("❌ 提取失败：检测到该视频没有声音轨道！")
                            else:
                                my_clip.audio.write_audiofile(temp_audio, logger=None)
                                my_clip.close()

                                # 3. 语音识别
                                model = load_whisper_model()
                                result = model.transcribe(
                                    temp_audio,
                                    language="zh",
                                    initial_prompt="以下是一段标准的简体中文普通话录音。",
                                    fp16=False
                                )

                                st.session_state['auto_text'] = result["text"]
                                st.success("✅ 提取成功！")
                                time.sleep(1)  # 停留1秒让用户看到成功提示
                                st.rerun()

                        except Exception as e:
                            # 【新增防护】无论出什么错，直接弹在网页上！
                            st.error(f"❌ 提取失败，底层报错信息：{str(e)}")
                            st.info(
                                "💡 提示：如果看到 'ffprobe' 或 'ffmpeg' 等字眼，请确保系统已通过 conda 安装了 ffmpeg。")

                        finally:
                            # 【新增防护】哪怕中途崩溃，也要把占硬盘的临时文件删掉
                            if os.path.exists(temp_vid):
                                try:
                                    os.remove(temp_vid)
                                except:
                                    pass
                            if os.path.exists(temp_audio):
                                try:
                                    os.remove(temp_audio)
                                except:
                                    pass
            else:
                st.image(multi_media, use_container_width=True)

    with col_text:
        multi_txt = st.text_area(
            "📝 第二步：输入配套文本（或点击左侧自动提取）",
            value=st.session_state['auto_text'],
            height=200,
            key="multi_txt"
        )

    # 联合检测按钮
    if multi_media and multi_txt:
        if st.button("🔍 启动跨模态融合鉴别", type="primary", use_container_width=True):
            st.divider()
            st.subheader("📊 联合检测报告")

            from text_module import analyze_text

            with st.spinner("并行调用视觉鉴别引擎与文本鉴别引擎..."):
                file_type = multi_media.name.split('.')[-1].lower()
                p_media = 0.0
                media_label = ""

                if file_type in ['mp4', 'avi', 'mov']:
                    from video_module import analyze_video
                    temp_path = f"temp_multi_video.{file_type}"
                    with open(temp_path, "wb") as f:
                        multi_media.seek(0) # 核心修复：倒带文件指针
                        f.write(multi_media.read())
                    res_media = analyze_video(temp_path)
                    p_media = res_media.get('avg_probability', 0)
                    media_label = "🎞️ 视频维度AI生成概率"
                    if os.path.exists(temp_path):
                        os.remove(temp_path)
                else:
                    from image_module import analyze_image
                    res_media = analyze_image(multi_media)
                    p_media = res_media['final_probability']
                    media_label = "🖼️ 图像维度AI生成概率"

                # 文本检测
                res_txt = analyze_text(multi_txt)
                p_txt = res_txt['final_probability']

                # 融合权重（视觉60%，文本40%）
                joint_prob = (p_media * 0.60) + (p_txt * 0.40)

            # 展示三指标
            col_d1, col_d2, col_d3 = st.columns(3)
            col_d1.metric(label=media_label, value=f"{p_media * 100:.1f}%")
            col_d2.metric(label="📝 文本维度AI生成概率", value=f"{p_txt * 100:.1f}%")
            col_d3.metric(
                label="🔗 多模态综合判定概率",
                value=f"{joint_prob * 100:.1f}%",
                delta="加权融合 (视觉0.6 + 文本0.4)",
                delta_color="off"
            )

            st.progress(float(joint_prob))

            # 综合判定
            if joint_prob >= high_risk_threshold:
                st.error("🚨 **联合研判结论：高度疑似AI生成的多模态内容**（图像/视频与文本均呈现显著生成特征）")
            elif warning_threshold <= joint_prob < high_risk_threshold:
                st.warning("⚠️ **联合研判结论：内容存疑**，可能存在局部AI修改或跨模态不一致，建议人工复核")
            else:
                st.success("✅ **联合研判结论：内容安全**，未见明显多模态伪造痕迹")