import sys import os import torch import torch.nn.functional as F import cv2 import numpy as np import argparse from PIL import Image from torchvision import transforms from tqdm import tqdm # === 配置 === # 使用 DINOv2 Giant 带寄存器版本 (修复背景伪影，最强版本) # 既然你不在乎显存，我们直接上 1.1B 参数的模型 MODEL_NAME = 'dinov2_vitg14_reg' DEVICE = "cuda" if torch.cuda.is_available() else "cpu" def init_model(): print(f"🚀 [系统] 初始化 DINOv2 ({MODEL_NAME})...") if DEVICE == "cuda": print(f"✅ [硬件确认] 正在使用显卡: {torch.cuda.get_device_name(0)}") print(f" (显存状态: {torch.cuda.memory_allocated()/1024**2:.2f}MB 已用)") else: print("❌ [警告] 未检测到显卡，Giant 模型在 CPU 上会非常慢！") # 加载模型 # force_reload=False 避免每次都下载 #model = torch.hub.load('facebookresearch/dinov2', MODEL_NAME) local_path = '/root/.cache/torch/hub/facebookresearch_dinov2_main' print(f"📂 [系统] 正在从本地缓存加载代码: {local_path}") if os.path.exists(local_path): model = torch.hub.load(local_path, MODEL_NAME, source='local') else: # 如果万一路径不对，再回退到在线加载（虽然大概率会失败） print("⚠️ 本地缓存未找到，尝试在线加载...") model = torch.hub.load('facebookresearch/dinov2', MODEL_NAME) model.to(DEVICE) model.eval() return model def preprocess_for_dino(img_cv): """ DINOv2 专用预处理： 1. 尺寸必须是 14 的倍数 2. 标准 ImageNet 归一化 """ h, w = img_cv.shape[:2] # 向下取整到 14 的倍数 new_h = (h // 14) * 14 new_w = (w // 14) * 14 img_resized = cv2.resize(img_cv, (new_w, new_h)) img_pil = Image.fromarray(cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) return transform(img_pil).unsqueeze(0).to(DEVICE), new_h, new_w def scan_and_draw(model, t1_path, t2_path, output_path, threshold): # 1. OpenCV 读取 img1_cv = cv2.imread(t1_path) img2_cv = cv2.imread(t2_path) if img1_cv is None or img2_cv is None: print("❌ 错误: 无法读取图片") return # 强制 Resize 对齐 (以现状图 T2 为准，逻辑保持不变) h_orig, w_orig = img2_cv.shape[:2] img1_cv = cv2.resize(img1_cv, (w_orig, h_orig)) print(f"🔪 [处理] DINOv2 扫描... 原始尺寸: {w_orig}x{h_orig}") # 2. 预处理 (DINO 需要整图输入，不再需要 sliding window 切片循环) # 但为了兼容 DINO 的 Patch 机制，我们需要微调尺寸为 14 的倍数 t1_tensor, h_align, w_align = preprocess_for_dino(img1_cv) t2_tensor, _, _ = preprocess_for_dino(img2_cv) print(f"🧠 [推理] Giant Model 计算中 (Patch网格: {h_align//14}x{w_align//14})...") with torch.no_grad(): # DINOv2 前向传播 (提取 Patch Token) # feat 形状: [1, N_patches, 1536] (Giant 的维度是 1536) feat1 = model.forward_features(t1_tensor)["x_norm_patchtokens"] feat2 = model.forward_features(t2_tensor)["x_norm_patchtokens"] # 计算余弦相似度 similarity = F.cosine_similarity(feat1, feat2, dim=-1) # [1, N_patches] # 3. 生成热力图数据 # reshape 回二维网格 grid_h, grid_w = h_align // 14, w_align // 14 sim_map = similarity.reshape(grid_h, grid_w).cpu().numpy() # 转换逻辑：相似度 -> 差异度 (Diff = 1 - Sim) heatmap_raw = 1.0 - sim_map # 将 14x14 的小格子放大回原图尺寸，以便与原图叠加 heatmap_avg = cv2.resize(heatmap_raw, (w_orig, h_orig), interpolation=cv2.INTER_CUBIC) # 统计信息 (逻辑保持不变) min_v, max_v = heatmap_avg.min(), heatmap_avg.max() print(f"\n📊 [统计] 差异分布: Min={min_v:.4f} | Max={max_v:.4f} | Mean={heatmap_avg.mean():.4f}") # ========================================== # 🔥 关键：保存原始灰度图 (逻辑保持不变) # ========================================== raw_norm = (heatmap_avg - min_v) / (max_v - min_v + 1e-6) cv2.imwrite("debug_raw_heatmap.png", (raw_norm * 255).astype(np.uint8)) print(f"💾 [调试] 原始热力图已保存: debug_raw_heatmap.png") # ========================================== # 5. 可视化后处理 (逻辑保持不变) # ========================================== # 归一化 (DINO 的差异通常在 0~1 之间，这里做动态拉伸以增强显示) # 如果差异非常小，max_v 可能很小，这里设置一个最小分母防止噪点放大 norm_factor = max(max_v, 0.4) heatmap_vis = (heatmap_avg / norm_factor * 255).clip(0, 255).astype(np.uint8) # 色彩映射 heatmap_color = cv2.applyColorMap(heatmap_vis, cv2.COLORMAP_JET) # 图像叠加 alpha = 0.4 blended_img = cv2.addWeighted(img2_cv, alpha, heatmap_color, 1.0 - alpha, 0) # 阈值过滤与画框 (逻辑完全保持不变) _, thresh_img = cv2.threshold(heatmap_vis, int(255 * threshold), 255, cv2.THRESH_BINARY) contours, _ = cv2.findContours(thresh_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) result_img = blended_img.copy() box_count = 0 # 既然用了 Giant，我们可以更精细地设定最小面积 # 此处保持和你之前代码一致的逻辑，但 DINO 不需要 PatchSize 参数，我们用原图比例 min_area = (w_orig * h_orig) * 0.005 # 0.5% 的面积 for cnt in contours: area = cv2.contourArea(cnt) if area > min_area: box_count += 1 x, y, bw, bh = cv2.boundingRect(cnt) # 画框 (白色粗框 + 红色细框) cv2.rectangle(result_img, (x, y), (x+bw, y+bh), (255, 255, 255), 4) cv2.rectangle(result_img, (x, y), (x+bw, y+bh), (0, 0, 255), 2) # 显示分数 # 计算该区域内的平均差异 region_score = heatmap_avg[y:y+bh, x:x+bw].mean() label = f"{region_score:.2f}" # 标签背景与文字 cv2.rectangle(result_img, (x, y-25), (x+80, y), (0,0,255), -1) cv2.putText(result_img, label, (x+5, y-7), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2) # 保存最终结果 cv2.imwrite(output_path, result_img) print("="*40) print(f"🎯 扫描完成! 发现区域: {box_count} 个") print(f"🖼️ 结果已保存至: {output_path}") print("="*40) if __name__ == "__main__": parser = argparse.ArgumentParser(description="DINOv2 Giant 违建热力图检测 (结构敏感版)") parser.add_argument("t1", help="基准图") parser.add_argument("t2", help="现状图") parser.add_argument("out", nargs="?", default="dino_result.jpg", help="输出图片名") # 为了兼容你的习惯，保留了 crop/step 参数接口，虽然 DINO 不需要它们 parser.add_argument("-c", "--crop", type=int, default=224, help="(已忽略) DINOv2 全图推理") parser.add_argument("-s", "--step", type=int, default=0, help="(已忽略) DINOv2 全图推理") parser.add_argument("-b", "--batch", type=int, default=16, help="(已忽略) DINOv2 全图推理") # 核心参数 # DINO 的 Cosine 差异通常比 DreamSim 小，建议阈值给低一点 (如 0.25 - 0.35) parser.add_argument("--thresh", type=float, default=0.30, help="检测阈值 (0.0-1.0)") args = parser.parse_args() # 初始化并运行 model = init_model() scan_and_draw(model, args.t1, args.t2, args.out, args.thresh)