UniVAD：一种免训练的少样本视觉异常检测统一模型（核心代码框架理解）

环境配置：python3.9

requriement: 全部安装依赖

本地上传：
pydensecrf @ git+https://github.com/lucasb-eyer/pydensecrf.git
ram @ git+https://github.com/xinyu1205/recognize-anything.git
segment_anything @ git+https://github.com/facebookresearch/segment-anything.git

数据集准备：MVTec AD

UniVAD核心代码理解框架文档

开源代码：https://uni-vad.github.io

第1章：代码全景与核心思想

1.1 项目定位

• 统一异常检测框架：同时支持像素级监督（MVTec）和图像级监督（医学数据）

• 少样本设定：k_shot=1 时仅依赖1张正常样本建模

• 三大对象类型：TEXTURE（纹理）、SINGLE（单一对象）、MULTI（多组件对象）

1.2 核心创新点代码映射

第2章：类结构与初始化 (__init__)

2.1 多模型加载

self.clip_model # CLIP (ViT-L-14-336) – 文本-图像对齐
self.dinov2_net # DINOv2 (vitg14) – 密集特征提取
self.dino_net # DINO (旧版) – 组件分解专用
self.cfa # CFA模块 – 特征聚合
self.decoder # LinearLayer – CLIP token降维

• 关键：三个视觉模型并行工作，特征维度分别为 1024 (CLIP), 1536 (DINOv2), 384 (DINO)

2.2 文本提示编码

self.text_prompts = encode_text_with_prompt_ensemble(...)
# 结果：{"object": tensor}，shape [2, 768] (正常/异常二分类)

2.3 数据变换管道

batch (原始图像)
├─→ self.transform_clip → CLIP输入 (336px, CLIP均值方差)
└─→ self.transform_dino → DINO输入 (336px, ImageNet均值方差)

• 注意：同一图像两种归一化，避免特征域冲突

第3章：前向传播 (forward) 深度解析

3.1 特征提取阶段 (Line 150-180)

# CLIP: 全局token + 7层中间层patch tokens
image_features, patch_tokens = self.clip_model.encode_image(..., out_layers=[6,12,18,24])
# DINOv2: 最后一层patch tokens
dino_patch_tokens = self.dinov2_net.forward_features(...)["x_norm_patchtokens"]

• 输出维度：

  • image_features: [B, 768] (全局)

  • patch_tokens: List[7×[B, 576, 1024]] (7层特征，每层576个14×14 token)

  • dino_patch_tokens: [B, 576, 1536]

3.2 全局异常分数 (Line 182-188)

global_score = 1 – (image_features @ normal_image_features.T).max()

• 物理意义：测试图像与正常原型的最大余弦相似度越低，越异常

3.3 Patch级异常图生成 (Line 191-224)

3.3.1 CLIP分支

• 原理：每个patch token vs 所有正常patch tokens，取最大相似度

• 代码模式：

  for i in range(len(patch_tokens)):
  if i % 2 == 0: continue # 只用奇数层（降采样减少计算）
  patch_tokens_reshaped = patch_tokens[i].view(576, 1, 1024)
  normal_tokens_reshaped = normal_patch_tokens[i].reshape(1, –1, 1024)
  sim_max = cosine_similarity(...).max(dim=1) # 576个patch各自的最相似原型

• 输出：anomaly_map_ret [1, 1, 336, 336]，值域[0,1]（异常概率）

3.3.2 DINOv2分支

• 原理：同CLIP，但特征维度1536，计算anomaly_map_ret_dino

3.3.3 CLIP-文本对齐分支

• 原理：patch tokens @ text_prompts，用softmax概率差作为异常图

• 关键代码：

  anomaly_map_vl = 100.0 * patch_tokens[layer] @ text_features # 温度缩放
  anomaly_map_vl = (softmax[:,1,:,:] – softmax[:,0,:,:] + 1) / 2 # 后处理

3.4 对象类型门控分支 (Line 235-520)

3.4.1 TEXTURE 分支 (Line 237-249)

• 适用：大面积纹理（如背景、匀质表面）

• 策略：全局特征融合，无部件分解

• 公式：anomaly_map = (CLIP + DINOv2 + text) / 3

3.4.2 SINGLE 分支 (Line 251-345)

• 适用：单一对象（如单个螺丝）

• 关键：用SAM mask裁剪对象区域，只在mask内计算patch相似度

• 缓存机制：

  self.normal_clip_part_patch_features[layer][0] # 单对象的正常原型

3.4.3 MULTI 分支 (Line 347-520) 最复杂

• 适用：多组件对象（如早餐盒含多个物品）

• 四步走：

Step 1: 组件分解 (get_heatmaps)

# 1. DINO特征聚类 → K-means生成粗粒度组件
# 2. SAM mask作为细粒度真值
# 3. assign_fine_to_coarse_torch: 将聚类结果对齐到SAM mask

• 核心函数：assign_fine_to_coarse_torch – 基于IoU的mask合并算法

Step 2: 部件级异常图 (Line 383-462)

• 缓存：self.normal_dino_part_patch_features[component_idx] 每个组件独立原型

• 关键：thresh[thresh>0] 只保留当前组件的patches

Step 3: 组件特征提取 (Line 464-508)

• 几何特征：面积、颜色直方图、中心坐标

• 视觉特征：CLIP/DINOv2组件级特征（经CFA聚合）

• 代码：

  self.component_feature_extractor.extract(image, query_masks) # 返回dict

Step 4: 组件对比 (Line 487-508)

sim_clip_component = cosine_similarity(query_clip_feat, all_normal_clip_feats)
sim_dino_component = cosine_similarity(query_dino_feat, all_normal_dino_feats)
sim_geo = cosine_similarity(query_geo_feat, all_normal_geo_feats)
anomaly_map_dist += (1 – sim.max()) # 组件级异常累加到像素图

第4章：Setup 方法——少样本原型构建

4.1 输入数据结构

data = {
"few_shot_samples": tensor[B_shot, 3, H, W], # k_shot张正常图
"dataset_category": str, # 类别名（用于缓存路径）
"image_path": [str, ...] # 对应文件路径
}

4.2 门控决策逻辑 (Line 588-604)

object_ratio > 0.65 and len(grounded_sam_masks[0]) <= 2 → TEXTURE
len(grounded_sam_masks[0]) == 1 → SINGLE
else → MULTI

• 物理意义：SAM分割出的mask数量+面积占比判断对象复杂度

4.3 原型特征缓存机制

4.4 MULTI类型训练时组件分解 (Line 629-720)

• K-means聚类：在DINO特征空间中找到部件原型（train_features_sampled）

• 贪心采样：GreedyCoresetSampler 保留代表性特征点

• 迭代优化：通过 while part_num not in part_num_right 动态调整聚类数

• 缓存：将聚类中心存为 ./heat_masks/{class_name}_heat/train_features_sampled.pth

第5章：关键辅助函数与技巧

5.1 get_heatmaps 函数 (Line 45-60)

• 跨注意力计算：query_feature @ dino_feature → 组件激活图

• CRF后处理：dense_crf 平滑边界（DINO+DINOv2双路）

5.2 assign_fine_to_coarse_torch 函数 (Line 732-759)

• 功能：将K-means的细粒度聚类mask对齐到SAM的粗粒度部件mask

• 算法：基于IoU最大匹配，循环分配每个细mask到最重叠的粗mask

5.3 CFA 类 (Line 763-777)

• 功能：特征上下文聚合（Context Feature Aggregation）

• 公式：A = softmax(Similarity), output = A @ V

• 用途：对组件级特征进行自适应加权聚合

第6章：数据流与调试技巧

6.1 输入 → 输出完整路径

输入图像 (336×336)
├─→ Transform → CLIP/DINOv2双路编码
├─→ Setup缓存 → 正常原型特征
├─→ Forward计算 → 三路异常图 (CLIP/DINOv2/Text)
├─→ Gate分支 → 根据对象类型融合
│ ├─→ TEXTURE: 直接平均
│ ├─→ SINGLE: SAM mask裁剪后平均
│ └─→ MULTI: 组件分解 + 四元特征对比
└─→ 输出: pred_score + pred_mask

6.2 调试断点建议

Python

复制

# 在 test_univad.py 中
model = UniVAD()
model.setup(data) # 断点1：看门控决策和缓存

output = model(batch, image_path)
# 断点2：检查 output['pred_mask'] 的数值范围
# 断点3：检查 global_score 与 mask 最大值关系

6.3 可视化验证

• 中间特征：在 forward() 中打印 anomaly_map_ret, anomaly_map_vls 的形状和均值

• 组件分解：在 setup() 的MULTI分支，保存 heatmap_refined 看K-means效果

• SAM mask：检查 ./masks/ 目录下的grounding_mask.png是否正确

第7章：性能优化与内存注意

7.1 计算瓶颈

• DINOv2 1536维特征：cosine_similarity 在576×N个patch上计算，内存占用高

• K-means：训练时fit过程在CPU，大数据集慢，建议用MiniBatchKMeans

7.2 缓存文件管理

• 路径：./heat_masks/{class_name}_heat/train/{shot_idx}/

• 关键文件：

  • grounding_mask.png (SAM生成)

  • heatresult_refined.png (K-means对齐结果)

  • train_features_sampled.pth (聚类中心)

7.3 超参敏感点

• self.out_layers = [6, 12, 18, 24]：CLIP中间层选择，影响patch token质量

• percentage=0.01：贪心采样比例，影响原型代表性

• n_cluster：K-means聚类数，决定组件分解粒度