基于 Stable Diffusion 的数据增强实践方案

负向 prompt 是隐藏宝藏：把'不想要的'写进去，比单纯堆正向词更能减少废图。

3. ControlNet：给'画家'一把尺子

纯文本容易'抽卡'，ControlNet 把 Canny 边缘、深度图、语义分割 mask 变成'草图'，让生成结果结构不变、纹理随便换。
做数据增强时，原图边缘图 + 随机 prompt = 同一结构不同外观，完美。

4. LoRA：不煮大锅饭，只开小灶

全量微调 4 GB 模型？老板不给显卡。
LoRA 把权重更新拆成两个小矩阵，训练量降到 1/10，10 张图 10 分钟就能学会'某种裂纹风格'，迁移学习神器。

搭一条'可控'的增强流水线

'没有流程的生成，都是玄学。'

下面给出一条工业界能落地的 Python 流水线，每一步都能打断点 debug，拒绝黑箱。

0. 环境一键复现

# 建议用 conda，别问，问就是省头发
conda create -n sdaug python=3.10
conda activate sdaug
pip install diffusers==0.21.0 transformers accelerate xformers opencv-python safetensors
# 显卡>=16G 可不开 xformers，<16G 建议加上，省显存

1. 原图→边缘图：保留结构

# canny_extract.py
import cv2
import os

def extract_canny(img_path, low=100, high=200, output_size=512):
    img = cv2.imread(img_path)
    img = cv2.resize(img, (output_size, output_size))
    canny = cv2.Canny(img, low, high)
    # 扩通道，适配 ControlNet 输入
    canny = cv2.cvtColor(canny, cv2.COLOR_GRAY2RGB)
    return canny

# 批量处理
os.makedirs("canny_dir", exist_ok=True)
for f in os.listdir("raw_images"):
    canny = extract_canny(f"raw_images/{f}")
    cv2.imwrite(f"canny_dir/{f}", canny)

小贴士：Canny 阈值别手抖，低阈值太高会把细节弄丢，经验值 50/150 起步，每张图都用同样阈值，保证后续对齐。

2. prompt 模板：把'随机'装进笼子里

# prompt_bank.py
templates = {
    "crack": [
        "a photo of {defect} on {object}, industrial scene, {lighting}, 4K, sharp, no text",
        "close-up shot of {defect} defect, metal surface, {lighting}, realistic, high contrast"
    ],
    "lighting": ["under factory LED light", "natural daylight", "dim warehouse light", "fluorescent tube light"],
    "object": ["aluminum panel", "steel plate", "phone screen", "car bumper"]
}

def sample_prompt(defect="crack"):
    import random
    t = random.choice(templates[defect])
    lighting = random.choice(templates["lighting"])
    obj = random.choice(templates["object"])
    return t.format(defect=defect, lighting=lighting, object=obj)

模板化 = 可复现 + 可单元测试，别小看这一步，后期排查语义漂移全靠它。

3. 图像→图像：把边缘图喂给 Stable Diffusion

# sd_aug.py
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import os

# 1. 加载 ControlNet
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
).to("cuda")

# 2. 内存优化三板斧
pipe.enable_xformers_memory_efficient_attention()
pipe.enable_attention_slicing()
pipe.enable_model_cpu_offload()
# batch 大时开

# 3. 批量生成
os.makedirs("aug_images", exist_ok=True)
for idx, canny_file in enumerate(os.listdir("canny_dir")):
    canny_image = load_image(f"canny_dir/{canny_file}")
    prompt = sample_prompt(defect="crack")
    negative = "cartoon, painting, lowres, blurry, extra fingers, text, watermark"
    out = pipe(
        prompt=prompt,
        negative_prompt=negative,
        image=canny_image,
        num_inference_steps=30,
        guidance_scale=7.5,
        generator=torch.Generator().manual_seed(42+idx),  # 可复现
        strength=0.9  # 0~1，越大越偏离原图
    ).images[0]
    out.save(f"aug_images/{idx:04d}.jpg")

strength 参数是灵魂：0.7 以下：基本只是'重新打光'；0.9 左右：结构保留但纹理大换血；1.0：放飞自我，可能把裂纹画成涂鸦。

4. 自动过滤：别让'垃圾'进数据集

# filter.py
from transformers import CLIPProcessor, CLIPModel
import torch, os, json
from PIL import Image

clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to("cuda")
proc = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

def clip_score(image, text):
    inputs = proc(text=[text], images=image, return_tensors="pt").to("cuda")
    with torch.no_grad():
        outputs = clip(**inputs)
        logits = outputs.logits_per_image
    return logits.item()

threshold = 28  # 经验值，按业务调
manifest = []
for imgf in os.listdir("aug_images"):
    img = Image.open(f"aug_images/{imgf}")
    score = clip_score(img, "a photo of cracked phone screen")
    if score >= threshold:
        manifest.append({"file": imgf, "clip": score})

json.dump(manifest, open("valid_images.json", "w"), ensure_ascii=False, indent=2)
print(f"过滤后剩余 {len(manifest)} 张，淘汰率 {1-len(manifest)/len(os.listdir('aug_images')):.2%}")

CLIP 分数不是圣旨，只能当'初筛'。后续还要让业务分类器回测，看加入这些图后 validation 涨不涨，再决定要不要下调阈值。

三个真实到'掉头发'的落地案例

1. 医疗影像：给罕见病灶'加戏'

背景：某三甲放射科，早期肺结节 CT 只有 87 张阳性，阴性 3000+，模型快把阳性当成'外星人'。

方案：

• 用 nnUNet 把原图结节 mask 抠出来→生成深度图→喂 ControlNet；
• prompt 模板里加入'-1mm 薄层、低剂量、40 岁患者'等医学关键词，让生成图自带 CT 质感；
• 生成 800 张后，请主任肉眼筛掉 120 张'不像病灶'的废图（医生眼光毒辣，一眼看穿伪影）。

结果：

• 召回率从 0.61 → 0.78，假阳性降 35%；
• 论文投 MICCAI，审稿人唯一质疑：'伦理批准呢？'——生成数据也要补伦理批件，别踩坑。

2. 工业质检：把'缺陷'搬到不同产线

背景：手机盖板玻璃，裂纹样本 214 张，客户要求识别 5 条产线、3 种光照、4 种角度共 60 种工况。

方案：

• 用 Blender 批量渲染虚拟边缘图（裂纹形状固定，角度/光照随意调），一天造 2W 张边缘图；
• 边缘图 + ControlNet 生成真实纹理，**prompt 里随机