import streamlit as st import subprocess import os from PIL import Image # === 页面基础配置 === st.set_page_config(layout="wide", page_title="多核违建检测平台 v3.0") st.title("🚁 多核无人机违建检测平台 v3.0") st.caption("集成内核: DINOv2 Giant | DiffSim-Pro v2.1 | DreamSim Ensemble") # ========================================== # 1. 侧边栏：算法核心选择 # ========================================== st.sidebar.header("🧠 算法内核 (Core)") algo_type = st.sidebar.radio( "选择检测模式", ( "DreamSim(语义/感知) 🔥", #"DINOv2 Giant (切片/精度)", #"DiffSim (结构/抗视差)", #"DINOv2 Giant (全图/感知)" ), index=0, # 默认选中 DINOv2 切片版，因为它是目前效果最好的 help="DINOv2: 几何结构最敏感，抗光照干扰。\nDiffSim: 可调参数多，适合微调。\nDreamSim: 关注整体风格差异。" ) st.sidebar.markdown("---") st.sidebar.header("🛠️ 参数配置") # 初始化默认参数变量 script_name = "" cmd_extra_args = [] show_slice_params = True # 默认显示切片参数 # ========================================== # 2. 根据选择配置参数 # ========================================== if "DINOv2" in algo_type: # === DINOv2 通用配置 === st.sidebar.info(f"当前内核: {algo_type.split(' ')[0]}") # 阈值控制 (DINO 的差异值通常较小，默认给 0.3) thresh = st.sidebar.slider("敏感度阈值 (Thresh)", 0.0, 1.0, 0.30, 0.01, help="值越小越敏感，值越大越只看显著变化") # 区分全图 vs 切片 if "切片" in algo_type: script_name = "main_dinov2_sliced.py" show_slice_params = True st.sidebar.caption("✅ 适用场景：4K/8K 大图，寻找细微违建。") else: script_name = "main_dinov2.py" # 对应你之前的 main_dinov2_giant.py (全图版) show_slice_params = False # 全图模式不需要调切片 st.sidebar.caption("⚡ 适用场景：快速扫描，显存充足，寻找大面积变化。") elif "DiffSim" in algo_type: # === DiffSim 配置 === script_name = "main_finally.py" show_slice_params = True # DiffSim 需要切片 st.sidebar.subheader("1. 感知权重") w_struct = st.sidebar.slider("结构权重 (Struct)", 0.0, 1.0, 0.3, 0.05) w_sem = st.sidebar.slider("语义权重 (Sem)", 0.0, 1.0, 0.7, 0.05) w_tex = st.sidebar.slider("纹理权重 (Texture)", 0.0, 1.0, 0.0, 0.05) st.sidebar.subheader("2. 信号处理") kernel = st.sidebar.number_input("抗视差窗口 (Kernel)", value=5, step=2) gamma = st.sidebar.slider("Gamma 压制", 0.5, 4.0, 1.0, 0.1) thresh = st.sidebar.slider("可视化阈值", 0.0, 1.0, 0.15, 0.01) # 封装 DiffSim 独有的参数 cmd_extra_args = [ "--model", "Manojb/stable-diffusion-2-1-base", "--w_struct", str(w_struct), "--w_sem", str(w_sem), "--w_tex", str(w_tex), "--gamma", str(gamma), "--kernel", str(kernel) ] else: # DreamSim # === DreamSim 配置 === script_name = "main_dreamsim.py" show_slice_params = True st.sidebar.subheader("阈值控制") thresh = st.sidebar.slider("可视化阈值 (Thresh)", 0.0, 1.0, 0.3, 0.01) # ========================================== # 3. 公共参数 (切片策略) # ========================================== if show_slice_params: st.sidebar.subheader("🚀 扫描策略") # DINOv2 切片版建议 Batch 小一点 default_batch = 8 if "DINOv2" in algo_type else 16 crop_size = st.sidebar.number_input("切片大小 (Crop)", value=224) step_size = st.sidebar.number_input("步长 (Step, 0=自动)", value=0) batch_size = st.sidebar.number_input("批次大小 (Batch)", value=default_batch) else: # 全图模式，隐藏参数但保留变量防报错 st.sidebar.success("全图模式：无需切片设置 (One-Shot)") crop_size, step_size, batch_size = 224, 0, 1 # ========================================== # 4. 主界面：图片上传与执行 # ========================================== col1, col2 = st.columns(2) with col1: file_t1 = st.file_uploader("上传基准图 (Base / Old)", type=["jpg","png","jpeg"], key="t1") if file_t1: st.image(file_t1, use_column_width=True) with col2: file_t2 = st.file_uploader("上传现状图 (Current / New)", type=["jpg","png","jpeg"], key="t2") if file_t2: st.image(file_t2, use_column_width=True) st.markdown("---") # 启动按钮 if st.button("🚀 启动检测内核", type="primary", use_container_width=True): if not file_t1 or not file_t2: st.error("请先上传两张图片！") else: # 1. 保存临时文件 os.makedirs("temp_uploads", exist_ok=True) t1_path = os.path.join("temp_uploads", "t1.jpg") t2_path = os.path.join("temp_uploads", "t2.jpg") # 结果文件名根据算法区分，防止缓存混淆 result_name = f"result_{script_name.replace('.py', '')}.jpg" out_path = os.path.join("temp_uploads", result_name) with open(t1_path, "wb") as f: f.write(file_t1.getbuffer()) with open(t2_path, "wb") as f: f.write(file_t2.getbuffer()) # 2. 构建命令 # 基础命令: python3 script.py t1 t2 out cmd = ["python3", script_name, t1_path, t2_path, out_path] # 添加通用参数 (所有脚本都兼容 -c -s -b --thresh 格式) # 注意：即使全图模式脚本忽略 -c -s，传进去也不会报错，保持逻辑简单 cmd.extend([ "--crop", str(crop_size), "--step", str(step_size), "--batch", str(batch_size), "--thresh", str(thresh) ]) # 添加特定算法参数 (DiffSim) if cmd_extra_args: cmd.extend(cmd_extra_args) # 3. 显示状态与运行 st.info(f"⏳ 正在调用内核: `{script_name}` ...") st.text(f"执行命令: {' '.join(cmd)}") # 方便调试 try: # 实时显示进度条可能比较难，这里用 spinner with st.spinner('AI 正在进行特征提取与比对... (DINO Giant 可能需要几秒钟)'): result = subprocess.run(cmd, capture_output=True, text=True) # 4. 结果处理 # 展开日志供查看 with st.expander("📄 查看内核运行日志", expanded=(result.returncode != 0)): if result.stdout: st.code(result.stdout, language="bash") if result.stderr: st.error(result.stderr) if result.returncode == 0: st.success(f"✅ 检测完成！耗时逻辑已结束。") # 结果展示区 r_col1, r_col2 = st.columns(2) with r_col1: # 尝试读取调试热力图 (如果有的话) if os.path.exists("debug_raw_heatmap.png"): st.image("debug_raw_heatmap.png", caption="🔍 原始差异热力图 (Debug)", use_column_width=True) else: st.warning("无调试热力图生成") with r_col2: if os.path.exists(out_path): # 使用 PIL 打开以强制刷新缓存 (Streamlit 有时会缓存同名图片) res_img = Image.open(out_path) st.image(res_img, caption=f"🎯 最终检测结果 ({algo_type})", use_column_width=True) else: st.error(f"❌ 未找到输出文件: {out_path}") else: st.error("❌ 内核运行出错，请检查上方日志。") except Exception as e: st.error(f"系统错误: {e}")