193 lines
7.7 KiB
Python
193 lines
7.7 KiB
Python
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}")
|