工业质检数据不平衡难题：用Stable Diffusion生成缺陷图像提升分割模型性能4.6%

1. 项目概述与核心挑战

在工业生产的质检环节，尤其是基于光学图像（包括可见光和热成像）的自动化检测中，监督式机器学习模型已经成为了不可或缺的工具。这些模型能够以极高的精度和速度，识别出产品表面的划痕、装配错误，或是像我们这次项目中的核心——通过热成像分析联合收割机内部部件的温度异常，从而定位潜在的机械故障。然而，一个长期困扰工业AI落地的“老大难”问题，就是数据不平衡。简单来说，生产线上绝大多数产品都是合格的，能拍到的“正常”部件图像成千上万；而真正有缺陷的、需要被检出的“坏件”却凤毛麟角。这种“好”样本远多于“坏”样本的极端不平衡数据集，直接喂给模型训练，结果往往是模型学会了“躺平”——它倾向于把所有样本都预测为“正常”，因为这样就能获得很高的整体准确率，但代价是完全丧失了检出缺陷的能力。

传统的应对策略，比如在损失函数上做文章（如Focal Loss），或者使用旋转、裁剪、颜色抖动这类经典的数据增强，我都尝试过。前者需要对超参数进行非常精细的调校，效果不稳定；后者则更像是“隔靴搔痒”，只能改变图像的表面特征（几何、颜色），无法创造出具有真实语义信息的“新缺陷”。例如，你无法通过旋转一张正常的轴承图片，就让它“长出”一道真实的裂纹。这就像试图用复印机复印一张白纸，期望得到一幅彩色油画一样不现实。

因此，当生成式人工智能（GenAI）浪潮袭来时，我立刻意识到这可能是一个破局点。GenAI的核心能力是“创造”，它通过学习海量数据的内在分布，能够生成前所未见但符合逻辑的新内容。那么，我们能否让它学会“缺陷”的视觉模式，然后批量地将正常图像“改造”成缺陷图像，从而低成本、高效地平衡我们的数据集呢？这就是我们整个项目的出发点。我们选择了当前最热门的两个图像生成模型——Stable Diffusion和CycleGAN，在一个具体的工业场景（联合收割机部件热成像分割）中，系统地验证了这一想法。最终，Stable Diffusion交出了一份令人惊喜的答卷：将分割模型的平均交并比（Mean IoU）提升了4.6个百分点。这个数字在工业质检的精度追求中，意义重大。

1.1 核心需求：从“数据饥渴”到“数据创造”

这个项目的核心需求非常明确：在缺陷样本极其稀缺（仅有253张带缺陷的热成像图，而正常样本有近千张）的前提下，显著提升语义分割模型对联合收割机内部部件（发动机、马达、粉碎机）的定位精度，为后续基于温度阈值的缺陷判定提供更准确的区域基础。

传统的解决路径无外乎两条：一是投入巨大成本，去收集更多真实的缺陷样本，这在生产中往往意味着要等待故障发生，或者人为制造故障，成本高、周期长，甚至不现实；二是依赖算法技巧，在模型层面做优化，但天花板明显。GenAI为我们开辟了第三条路：数据合成。我们不需要去“找”更多缺陷，而是尝试去“造”出逼真的缺陷。这背后的技术价值在于，它首次将数据扩充的维度，从低级的像素扰动，提升到了高级的语义生成层面。模型不再只是学习如何对现有图片做微小的形变，而是去理解“什么是缺陷”——在热成像中，缺陷往往表现为异常的高温区域。生成模型的任务，就是学会在正常部件的对应位置，“画”出符合物理规律的高温斑块。

2. 技术方案选型与原理剖析

面对数据不平衡，为什么最终锁定了Stable Diffusion和CycleGAN这两个模型？这背后是基于对问题本质和模型特性的深度考量。我们需要的是一个“图像到图像”的翻译模型，输入是正常部件图像，输出是带有模拟缺陷的对应图像。同时，由于我们的数据是专业的热成像图，与自然图像分布差异巨大，预训练模型必须经过针对性的微调（Fine-tuning）才能奏效。

2.1 Stable Diffusion：基于扩散模型的文生图与图生图

Stable Diffusion是当前开源社区最强大的文生图模型之一。它的核心是一个“去噪扩散”过程：先生成一张纯噪声图片，然后通过一个训练好的去噪模型，一步步去除噪声，最终得到清晰的图像。这个去噪过程可以由文本提示词（Prompt）来引导，从而控制生成内容。

为什么选择它？

1. 高质量与可控性：扩散模型生成的图像质量普遍高于早期的GAN，细节更丰富。更重要的是，它通过文本提示词提供了极强的语义控制能力。我们可以用“一台Variant A型号的联合收割机，其发动机存在缺陷”这样的描述，来精确指导生成过程。
2. 图生图（Image-to-Image）模式：这正是我们需要的。我们可以输入一张正常的部件图，并设置一个“去噪强度”（strength）参数。这个参数控制了输出结果在多大程度上保留输入图的构图和布局（低强度），又有多大程度遵循文本提示词去创造新内容（高强度）。这确保了生成的缺陷图像中，部件的位置、大小与原始图像基本对齐，这对于后续分割模型的训练至关重要——标签掩膜（Mask）可以复用或微调。
3. 强大的社区与生态：基于Stable Diffusion的微调技术（如DreamBooth、LoRA）非常成熟，便于我们在小规模专业数据集上对其进行定制。

我们的微调策略：原始的Stable Diffusion是在海量自然图像（如LAION数据集）上训练的，根本不认识“联合收割机内部热成像图”。因此，我们必须对其进行微调。我们采用了冻结文本编码器和图像编码器，只训练核心的UNet去噪模型的策略。这样既能利用模型预训练获得的世界知识，又能让它快速适应我们独特的图像分布。我们为每张训练图片配上了精心设计的提示词，例如“一张RGB图像，内容是一台完好的Variant A型联合收割机（ch machine）”。这里特意用“ch machine”替代了“combine harvester”，是因为模型预训练知识中的“联合收割机”都是外观图，与我们的内部视图无关，强行使用反而会造成干扰。

2.2 CycleGAN：无配对图像翻译的经典之作

CycleGAN是图像翻译领域的经典模型，它的最大特点是不需要成对的数据。也就是说，我们不需要有一张正常的发动机图和一张与之严格对应的、同一台发动机的缺陷图。我们只需要两个集合：一个“正常部件图像”集合（A域），和一个“缺陷部件图像”集合（B域）。CycleGAN会自己学习这两个域之间的映射关系。

为什么选择它作为对比？

1. 解决无配对数据的痛点：在工业场景中，获取严格配对的“正常-缺陷”图像对极其困难。CycleGAN的设定更符合现实条件。
2. 风格迁移的直观性：它的学习目标非常直观：让生成的缺陷图看起来像是来自缺陷域（B域），同时还能通过一个反向生成器变回原始的正常图（循环一致性损失）。这迫使模型捕捉缺陷的本质“风格”，比如热成像中高温区域的亮度、纹理特征。
3. 作为重要的基线对比：与需要文本引导的Stable Diffusion相比，CycleGAN是完全数据驱动的。对比两者的效果，能帮助我们理解在工业数据增强任务中，引入先验知识（文本提示）是否更有优势。

两者的核心差异与选择逻辑：

• 控制粒度：Stable Diffusion通过提示词可以实现细粒度控制（如指定哪个部件缺陷），而CycleGAN是整体风格转换，可能让图像中所有部件都“变亮”，无法精确指定缺陷位置。
• 数据需求：Stable Diffusion需要图像-文本对进行微调，对数据标注（写提示词）有一定要求；CycleGAN只需要两个图像集，标注成本更低。
• 生成原理：Stable Diffusion是概率扩散，生成过程更具创造性和多样性；CycleGAN是确定性映射，生成结果可能更稳定，但多样性稍弱。

在我们的实验中，同时评估这两者，就是为了在实际工业场景中，给出一份基于不同约束条件（有无文本标注、是否需要精确控制）下的技术选型参考。

注意：模型选型没有绝对的好坏，只有是否适合。如果你的缺陷类型单一、且难以用语言精确描述，但拥有一定数量的缺陷图片集，CycleGAN可能是更快捷的起点。如果你的缺陷类别多、需要精确对应部件、且能总结出清晰的文本描述，那么投入精力微调Stable Diffusion会带来更大的回报。

3. 实操流程：从数据准备到模型训练

理论很美好，但落地过程充满了细节。下面我将以Stable Diffusion为例，拆解整个从数据准备到生成图像的全流程，并分享其中踩过的坑和总结的经验。

3.1 数据准备与预处理

我们的原始数据是1200张联合收割机的侧视热成像图，其中缺陷图仅253张。图像包含了多种机器变体（19种）和缺陷组合（9种）。

第一步：构建微调训练集

1. 精选代表性样本：我们没有使用全部数据。对于缺陷图像，我们从每种缺陷变体中手动挑选最具代表性的8张，共40张。对于正常图像，从18种变体（剔除1种无对应缺陷的变体）中各选5张，共90张。手动挑选至关重要，目的是确保训练集能覆盖所有重要的视觉模式，避免模型学到噪声或特例。
2. 生成文本提示词：为每一张训练图像编写对应的提示词。这里我们设计了两种模板：
   • 模板V1（细粒度）：“一张{图像类型}图像，内容是一台{变体型号}的联合收割机（ch machine），其{部件名称}存在缺陷。” 例如：“a rgb image of a variant A ch machine with a defective main engine”。
   • 模板V2（粗粒度）：“一张{图像类型}图像，内容是一台存在缺陷的{变体型号}联合收割机（ch machine）。” 例如：“a rgb image of a defective variant A ch machine”。
   • 对于正常图像，统一使用：“一张{图像类型}图像，内容是一台完好的{变体型号}联合收割机（ch machine）。”
3. 图像格式：我们将每张原始温度数据同时保存为RGB伪彩图和灰度图，以探究哪种格式对生成模型更友好。

第二步：准备生成源数据集为了评估生成效果，我们需要另一组“未见过的”正常图像作为生成的输入源。我们从训练集中未使用的正常图像里，为每种变体再挑选最多5张，组成一个包含80张图像（各有RGB和灰度版本）的源数据集。

3.2 Stable Diffusion模型微调实战

我们使用基于Keras的实现进行微调。关键步骤如下：

1. 加载预训练模型：使用stabilityai/stable-diffusion-2-1基础模型。
2. 冻结部分参数：为了保持模型的通用知识并防止在小数据集上过拟合，我们冻结了文本编码器（Text Encoder）和图像编码器（VAE Encoder），只训练UNet去噪模型。
3. 两阶段训练策略：这是本项目的一个关键技巧。由于缺陷样本（40张）远少于正常样本（90张），如果直接混合训练，模型可能更倾向于学习生成正常图像。
   • 第一阶段：使用全部130张（正常+缺陷）图像进行训练，让模型先学会所有机器变体的基本结构和热成像风格。
   • 第二阶段：在第一阶段训练好的模型基础上，仅使用40张缺陷图像进行继续微调。这相当于给模型一个“强化课程”，让它更加专注于学习“缺陷”特征。
4. 训练超参数：学习率设置为1e-5，使用AdamW优化器，批量大小（Batch Size）根据GPU内存（12GB RTX 3060）设为1。训练轮数（Epoch）需要密切监控损失曲线，通常在几十到一百轮左右，当验证损失不再明显下降时即可停止，避免过拟合。

实操心得：微调GenAI模型时，监控生成效果比只看损失曲线更重要。每隔几个Epoch，就用固定的提示词和种子（Seed）生成几张图片，直观地看生成质量是否有提升、是否出现模式崩溃（生成的图片都一样）。我们的实验发现，两阶段训练能显著改善缺陷特征的生成质量，避免了缺陷特征被“淹没”在正常样本中。

3.3 图像生成与超参数调优

微调好模型后，就可以用它对源数据集（80张正常图）进行“图生图”转换了。这个过程有多个超参数需要调试，它们直接影响生成图像的质量和多样性：

• 提示词变体（Prompt Variant）：我们对比了V1（指定缺陷部件）和V2（仅说机器有缺陷）。
• 图像类型（Image Type）：选择RGB还是灰度图。
• 是否使用第二阶段微调（Fine-tuning）：即是否使用经过缺陷数据强化训练的模型。
• 无条件引导尺度（Unconditional Guidance Scale, UGS）：这个值控制生成结果服从文本提示词的程度。值越高，生成内容越贴近描述，但可能牺牲图像自然度和多样性。我们测试了7.5和13.0。
• 图像强度（Image Strength）：在图生图模式下，这个值决定初始噪声图像与输入原图的相似度。值越低（如0.3），输出图像与原图构图越像；值越高（如0.7），生成器有更多自由发挥空间，可能创造出与原图布局差异更大的图像，但也可能丢失部件位置信息。我们测试了0.3, 0.5, 0.7。
• 生成数量与步数：每张输入图生成2个变体，扩散步数设为标准的50步。

我们系统地组合了这些参数，生成了多组扩充数据集。例如，一组数据可能是：使用V1提示词（指定发动机缺陷）、RGB图像、经过两阶段微调的模型、UGS=7.5、Image Strength=0.3所生成的所有图像。

4. 实验评估与结果分析

评估环节是验证“生成数据是否真有用”的核心。我们设计了一个严谨的三阶段实验流程。

4.1 实验设计：层层递进的验证

1. 阶段一：生成数据质量初筛（Model 1）

   • 目的：检验生成的数据是否“像”真实数据。如果模型在虚假数据上训练后，在真实数据上完全失效，那说明生成质量太差。
   • 操作：用纯生成的数据训练一个SegFormer分割模型（Model 1），然后在真实的测试集上评估其性能（Mean IoU）。
   • 结果解读：如表3所示，最好的生成数据组合（使用RGB图、UGS=7.5、强度0.3）能让Model 1在真实测试集上达到约0.65的Mean IoU。虽然远低于用真实数据训练的0.8，但这证明生成数据并非随机噪声，它确实包含了可被分割模型学习的、与真实数据相关的特征。这个阶段帮我们筛选出了最佳生成参数组合。

2. 阶段二：混合数据有效性验证（Model 2 vs Model 3）

   • 目的：核心验证——加入生成数据能否提升在真实数据上训练的模型性能。
   • 操作：
     • Model 3（基线）：仅使用原始真实数据的训练集训练SegFormer。
     • Model 2（实验组）：使用原始真实训练集 + 阶段一筛选出的最佳生成数据，共同训练另一个SegFormer。
   • 关键对比：将Model 2和Model 3在相同的真实测试集上对比Mean IoU。
   • 核心结果：如表4所示，仅用真实数据训练的基线模型Mean IoU为0.800。当加入Stable Diffusion生成的数据后，模型性能提升至0.846，相对提升4.6%。加入CycleGAN数据后，提升至0.831（+3.1%）。这个提升在工业质检的精度竞赛中是非常显著的，很可能意味着漏检率或误检率降低了数个百分比。

4.2 结果深度剖析与洞见

定量结果告诉我们什么？

1. Stable Diffusion优于CycleGAN：4.6% > 3.1%。这表明，在我们这个需要精确对应部件生成缺陷的任务中，通过文本提示词提供的细粒度控制带来了优势。
2. RGB格式优于灰度图：在两种模型的生成质量筛选中，RGB数据训练出的模型都远好于灰度图。这可能是因为RGB伪彩色包含了更丰富的对比度和纹理信息，有助于模型学习。
3. 适中的控制强度最佳：对于Stable Diffusion，较低的图像强度（0.3）和适中的引导尺度（7.5）效果最好。这意味着我们需要生成图像在保持原图部件布局（低强度）的基础上，适度地引入缺陷特征（合适的引导尺度），而不是天马行空地重绘。

定性分析（看图说话）的发现： 观察生成的图像（如图4，图5），我们发现了模型行为的微妙之处：

• Stable Diffusion的“过度生成”：当使用V1提示词要求生成“有缺陷的粉碎机”时，生成的图像中粉碎机区域确实变亮了（模拟高温），但发动机区域有时也会不明原因地变亮。这说明模型虽然理解了“缺陷=变亮”，但未能完美地将该概念与文本指定的部件绑定。这可能是因为训练数据中不同部件缺陷的共现模式被模型学到了。
• CycleGAN的“全体亮化”：由于CycleGAN没有部件级别的控制，它学习到的是“缺陷域”的整体风格。因此，它生成的缺陷图像，往往倾向于让画面中的所有部件都不同程度地变亮，而不是在特定位置生成缺陷。这在缺陷定位任务中会引入噪声。
• 稀有缺陷的生成难题：对于“马达缺陷”这类在真实数据中只有单一变体的稀有缺陷，两种模型都学得不好，生成效果不佳。这再次印证了机器学习的那句老话：垃圾进，垃圾出（Garbage in, garbage out）。生成模型的能力上限，依然受限于训练数据的多样性和质量。

5. 避坑指南与实战建议

基于这次完整的项目实践，我总结了以下几条对于想在工业场景中应用GenAI进行数据增强的同行们至关重要的建议：

1. 数据质量是生命线，宁可少而精

• 坑：盲目使用所有数据，包括模糊、标注不准、特征不明显的样本进行微调。
• 建议：像我们一样，手动筛选最具代表性的图像构建微调集。确保每一张训练图都是“教科书级别”的范例。对于生成模型，10张高质量图片的效果可能远好于100张嘈杂的图片。
• 技巧：可以先用基础模型生成一批数据，观察哪些类型的原图生成效果最差，反过来检查和清洗你的训练集。

2. 提示词工程需要“说模型能听懂的话”

• 坑：使用过于复杂、抽象或包含模型不熟悉概念的提示词。例如，直接使用“combine harvester”导致模型生成外观图。
• 建议：
  • 解构对象：将复杂对象分解为模型可能认知的基本元素。我们用“ch machine”来指代一个模型不认识的新概念。
  • 描述状态而非术语：用“高温区域”、“异常明亮的部位”来描述缺陷，可能比直接用“裂纹”、“磨损”等专业术语更有效，因为热成像的视觉表现就是温度差异。
  • 进行A/B测试：像我们一样，系统比较不同提示词变体（V1 vs V2）的效果。记录下哪种描述方式生成的图像更符合预期。

3. 超参数调优需要系统性的网格搜索

• 坑：凭感觉随便设置strength和guidance_scale，然后抱怨生成效果不好。
• 建议：必须像我们这样，设计一个系统的实验矩阵（如表1）。尤其是Image Strength，它直接决定了生成图像与原始标签的对应关系。对于分割任务，通常需要较低的值（0.2-0.4）来保持空间布局；对于分类任务，可以尝试更高的值以获得更多样性。

4. 评估体系必须包含“代理任务”

• 坑：只用人眼评判生成图像“像不像”，或者直接混合数据训练看最终指标，耗时漫长且无法指导中间调优。
• 建议：一定要设立一个快速反馈的代理评估任务。我们的“阶段一”（用纯生成数据训练分割模型）就是一个完美的代理任务。它用一个下游模型的性能，量化地反映了生成数据的质量。这个指标比人眼判断更客观，且能快速迭代不同的生成参数组合。

5. 理解并接受生成模型的局限性

• 坑：期望生成模型能完美解决“零样本”或“极端少样本”问题。
• 建议：生成模型是“放大器”，不是“无中生有器”。如果真实数据中某种缺陷只有一两张图，模型很难学会其精髓并生成多样化的样本。对于稀有类别，传统的过采样（Oversampling）、代价敏感学习（Cost-sensitive learning）与生成式增强相结合，可能是更稳妥的方案。我们的项目也指出，未来需要为稀有缺陷补充更多真实数据，才能进一步提升生成质量。

6. 未来展望与应用扩展

这次成功的实验，为我们打开了一扇门。它验证了在专业的、数据稀缺的工业视觉领域，GenAI的数据增强能力是切实有效的。接下来的探索方向很明确：

1. 任务泛化：我们在语义分割任务上成功了，那么对于目标检测（需要生成带边界框的缺陷）和图像分类任务呢？这需要调整生成和训练流程，例如探索ControlNet等引入空间控制的条件生成模型，来生成带定位信息的缺陷。

2. 多模态融合：目前我们只用了热成像。在实际质检中，往往结合可见光、X光等多源信息。能否利用多模态大模型（如基于CLIP的模型），实现跨模态的条件生成？例如，用文本描述缺陷，同时生成对应的热成像和可见光异常图像。

3. 在线持续学习：生产线上的缺陷模式可能会随时间漂移（如设备老化导致缺陷形态变化）。能否建立一个闭环系统，将新检出的真实缺陷图像自动加入微调数据集，让生成模型持续进化，实现“越用越聪明”的在线数据增强？

4. 合成数据质量评估标准化：目前缺乏权威的、针对工业合成数据的评估标准。除了用下游任务性能做评估，是否需要开发新的指标，从“物理合理性”（生成的热斑是否符合热传导规律）、“多样性”等维度进行衡量？

这个项目让我深刻体会到，AI在工业领域的落地，正从“感知智能”走向“创造智能”。我们不再只是被动地分析数据，而是开始主动地创造高质量的数据，来喂养和提升我们的模型。Stable Diffusion那4.6%的性能提升，不仅仅是一个数字，它代表了一种新的生产力工具正在成熟，一条解决工业AI“数据荒”的新路径正在被趟平。对于每一位从事工业AI、特别是面临数据困境的工程师来说，现在正是拿起生成式AI这把“利器”的最佳时机。