通过类似数据蒸馏或自动学习采样的方法，更加高效地学习良品数据分布 ...

好的，我们先聚焦第一个突破点：
   通过类似数据蒸馏或自动学习采样的方法，更加高效地学习良品数据分布。
  这里我提供一个完备的代码示例：
✅ Masked图像重建 + 残差热力图

这属于自监督蒸馏方法的一个变体：

• 利用一个 预练习MAE模子（或轻量ViT）对正常样本进行遮挡重建
  
• 用重建图与原图的残差来反映“异常程度”
  

---

✅ 示例环境依靠

1. pip install timm einops torchvision matplotlib

复制代码

---

✅ 完备代码（以MVTec中的图像为例）

1. import torch
2. import torch.nn as nn
3. import torchvision.transforms as T
4. from torchvision.utils import save_image
5. from torchvision.datasets.folder import default_loader
6. from einops import rearrange
7. import timm
8. import matplotlib.pyplot as plt
9. import os
10. from glob import glob
11. from PIL import Image
12. import numpy as np
14. # ---------------------------
15. # 模型定义：ViT作为Encoder + 简单Decoder
16. # ---------------------------
17. class MAE(nn.Module):
18.     def __init__(self, encoder_name='vit_base_patch16_224', mask_ratio=0.4):
19.         super().__init__()
20.         self.encoder = timm.create_model(encoder_name, pretrained=True)
21.         self.mask_ratio = mask_ratio
22.         self.patch_size = self.encoder.patch_embed.patch_size[0]
23.         self.num_patches = self.encoder.patch_embed.num_patches
24.         self.embed_dim = self.encoder.embed_dim
25.         self.decoder = nn.Sequential(
26.             nn.Linear(self.embed_dim, self.embed_dim),
27.             nn.GELU(),
28.             nn.Linear(self.embed_dim, self.patch_size**2 * 3)
29.         )
31.     def forward(self, x):
32.         B, C, H, W = x.shape
33.         x_patch = self.encoder.patch_embed(x)  # [B, num_patches, dim]
34.         B, N, D = x_patch.shape
36.         # 随机遮挡
37.         rand_idx = torch.rand(B, N).argsort(dim=1)
38.         num_keep = int(N * (1 - self.mask_ratio))
39.         keep_idx = rand_idx[:, :num_keep]
41.         x_keep = torch.gather(x_patch, 1, keep_idx.unsqueeze(-1).expand(-1, -1, D))
42.         x_encoded = self.encoder.blocks(x_keep)
43.         x_decoded = self.decoder(x_encoded)
45.         # 恢复顺序（只对keep部分重建）
46.         output = torch.zeros(B, N, self.patch_size**2 * 3).to(x.device)
47.         output.scatter_(1, keep_idx.unsqueeze(-1).expand(-1, -1, self.patch_size**2 * 3), x_decoded)
48.         output = rearrange(output, 'b n (p c) -> b c (h p) (w p)',
49.                            p=self.patch_size, c=3, h=int(H/self.patch_size), w=int(W/self.patch_size))
50.         return output
52. # ---------------------------
53. # 数据加载 + 预处理
54. # ---------------------------
55. transform = T.Compose([
56.     T.Resize((224, 224)),
57.     T.ToTensor(),
58.     T.Normalize([0.5]*3, [0.5]*3)
59. ])
61. inv_transform = T.Compose([
62.     T.Normalize(mean=[-1]*3, std=[2]*3)
63. ])
65. def load_images(path):
66.     files = sorted(glob(os.path.join(path, '*.png')) + glob(os.path.join(path, '*.jpg')))
67.     images = []
68.     for f in files:
69.         img = default_loader(f)
70.         images.append(transform(img))
71.     return torch.stack(images)
73. # ---------------------------
74. # 测试图像 → 重建图像 → 残差热图
75. # ---------------------------
76. def visualize_anomaly(original, recon, save_path='result.png'):
77.     residual = (original - recon).abs().sum(dim=1, keepdim=True)
78.     residual = residual / residual.max()
80.     fig, axs = plt.subplots(1, 3, figsize=(12, 4))
81.     axs[0].imshow(inv_transform(original[0]).permute(1, 2, 0).cpu().numpy())
82.     axs[0].set_title('Original')
83.     axs[1].imshow(inv_transform(recon[0]).permute(1, 2, 0).cpu().numpy())
84.     axs[1].set_title('Reconstruction')
85.     axs[2].imshow(residual[0, 0].cpu().numpy(), cmap='hot')
86.     axs[2].set_title('Anomaly Map')
87.     for ax in axs: ax.axis('off')
88.     plt.tight_layout()
89.     plt.savefig(save_path)
90.     plt.close()
92. # ---------------------------
93. # 主程序执行
94. # ---------------------------
95. if __name__ == '__main__':
96.     device = 'cuda' if torch.cuda.is_available() else 'cpu'
97.     model = MAE().to(device)
98.     model.eval()
100.     # 替换为 MVTec / VisA 任一类别路径
101.     image_dir = './mvtec/bottle/good/'  # 只加载良品图像
102.     images = load_images(image_dir).to(device)
104.     with torch.no_grad():
105.         for i in range(min(5, len(images))):
106.             input_img = images[i:i+1]
107.             recon_img = model(input_img)
108.             visualize_anomaly(input_img, recon_img, f'output_{i}.png')

复制代码

---

✅ 示例输出（保存为output_0.png等）：

• 左：原图
  
• 中：重建图（模子“明白的良品”）
  
• 右：异常热图（残差）
  

在正常样本上，残差图应接近0；如果输入的是异常图像，则对应区域将出现高相应。

---

✅ 可扩展方向

模块可扩展优化Encoder更换为轻量ViT（如 vit_tiny_patch16_224）Mask计谋利用结构化遮挡（如Block Mask）提拔重建挑战异常图像输入异常样本（如MVTec测试集中defect图）验证泛化本领练习到场重建loss微调，进步良品建模精度

---

如果你希望我进一步扩展为：

• 支持少量异常图像的快速修正版本；
  
• 或到场自动样本选择机制；
  

很好，我们继续在上一套基于自监督重建（MAE）方法的基础上，
为其 封装 Gradio Demo，以实现更直观的异常检测体验。

---

✅ 新增功能目标

• 上传恣意图片（良品或异常图）
  
• 及时体现：
  
  
  
  • 原图
    
  • 模子重建图
    
  • 残差热力图（高相应 = 异常区域）
    
  
  

---

✅ 完备代码（附Gradio界面）

1. import torch
2. import torch.nn as nn
3. import torchvision.transforms as T
4. from torchvision.utils import save_image
5. from torchvision.datasets.folder import default_loader
6. from einops import rearrange
7. import timm
8. import gradio as gr
9. import numpy as np
10. from PIL import Image
11. import matplotlib.pyplot as plt
12. import io
14. # ---------------------------
15. # 模型定义（同上）
16. # ---------------------------
17. class MAE(nn.Module):
18.     def __init__(self, encoder_name='vit_base_patch16_224', mask_ratio=0.4):
19.         super().__init__()
20.         self.encoder = timm.create_model(encoder_name, pretrained=True)
21.         self.mask_ratio = mask_ratio
22.         self.patch_size = self.encoder.patch_embed.patch_size[0]
23.         self.num_patches = self.encoder.patch_embed.num_patches
24.         self.embed_dim = self.encoder.embed_dim
25.         self.decoder = nn.Sequential(
26.             nn.Linear(self.embed_dim, self.embed_dim),
27.             nn.GELU(),
28.             nn.Linear(self.embed_dim, self.patch_size**2 * 3)
29.         )
31.     def forward(self, x):
32.         B, C, H, W = x.shape
33.         x_patch = self.encoder.patch_embed(x)
34.         B, N, D = x_patch.shape
35.         rand_idx = torch.rand(B, N).argsort(dim=1)
36.         num_keep = int(N * (1 - self.mask_ratio))
37.         keep_idx = rand_idx[:, :num_keep]
39.         x_keep = torch.gather(x_patch, 1, keep_idx.unsqueeze(-1).expand(-1, -1, D))
40.         x_encoded = self.encoder.blocks(x_keep)
41.         x_decoded = self.decoder(x_encoded)
43.         output = torch.zeros(B, N, self.patch_size**2 * 3).to(x.device)
44.         output.scatter_(1, keep_idx.unsqueeze(-1).expand(-1, -1, self.patch_size**2 * 3), x_decoded)
45.         output = rearrange(output, 'b n (p c) -> b c (h p) (w p)',
46.                            p=self.patch_size, c=3, h=int(H/self.patch_size), w=int(W/self.patch_size))
47.         return output
49. # ---------------------------
50. # 预处理 & 后处理
51. # ---------------------------
52. transform = T.Compose([
53.     T.Resize((224, 224)),
54.     T.ToTensor(),
55.     T.Normalize([0.5]*3, [0.5]*3)
56. ])
58. inv_transform = T.Compose([
59.     T.Normalize(mean=[-1]*3, std=[2]*3)
60. ])
62. def tensor_to_pil(t):
63.     t = inv_transform(t.squeeze(0)).clamp(0, 1)
64.     return T.ToPILImage()(t)
66. def residual_map(orig, recon):
67.     residual = (orig - recon).abs().sum(dim=1, keepdim=True)
68.     residual = residual / (residual.max() + 1e-8)
69.     heat = residual.squeeze().cpu().numpy()
70.     fig, ax = plt.subplots()
71.     ax.imshow(heat, cmap='hot')
72.     ax.axis('off')
73.     buf = io.BytesIO()
74.     plt.savefig(buf, format='png')
75.     plt.close(fig)
76.     buf.seek(0)
77.     return Image.open(buf)
79. # ---------------------------
80. # 推理函数
81. # ---------------------------
82. device = 'cuda' if torch.cuda.is_available() else 'cpu'
83. model = MAE().to(device)
84. model.eval()
86. def infer(img_pil):
87.     img_tensor = transform(img_pil).unsqueeze(0).to(device)
88.     with torch.no_grad():
89.         recon = model(img_tensor)
90.     recon_img = tensor_to_pil(recon)
91.     input_img = tensor_to_pil(img_tensor)
92.     heatmap = residual_map(img_tensor, recon)
93.     return input_img, recon_img, heatmap
95. # ---------------------------
96. # Gradio UI
97. # ---------------------------
98. demo = gr.Interface(
99.     fn=infer,
100.     inputs=gr.Image(type="pil", label="上传图像"),
101.     outputs=[
102.         gr.Image(type="pil", label="原图"),
103.         gr.Image(type="pil", label="重建图"),
104.         gr.Image(type="pil", label="残差热图")
105.     ],
106.     title="基于良品数据的异常检测（MAE重建）",
107.     description="上传图像，模型将重建正常区域并生成异常残差热力图"
108. )
110. if __name__ == '__main__':
111.     demo.launch()

复制代码

---

✅ 利用效果

你可以上传如下类型图像进行及时检测：

• ✔️ 良品图像：残差图团体应较为平滑，相应值低；
  
• ❌ 异常图像（如划痕/破损）：残差图中异常区域显着发亮（高相应）；
  

---

✅ 后续扩展建议：

模块可增强重建网络替换为 DRAEM / Reverse Distillation异常评分盘算全图平均残差 + Otsu二值化分割多样本比较支持目录上传并批量可视化迁移微调用少量目标数据 fine-tune 提拔领域鲁棒性

---

需要我下一步为你实现：

• ✅ 残差异常评分 + 二值掩码输出？
  
• ✅ 支持少量异常样本微调功能？
  
• ✅ 用 PatchCore / AnomalyCLIP 替换 MAE 结构？
  

你可以指定下一个要增强的方向，我这边可以直接给出代码。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。