YOLO+OpenClaw+AIGC缺陷生成：破解工业检测“数据荒”的低代码方案

import cv2
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image

# 加载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")

def generate_defect(good_image_path, mask_path, prompt):
    """
    在良品图上生成缺陷
    - good_image_path: 良品图像路径
    - mask_path: 缺陷区域掩码（白色区域表示要生成缺陷的位置）
    - prompt: 缺陷描述，如 "a deep scratch on metal surface"
    """
    image = load_image(good_image_path)
    mask = load_image(mask_path)
    
    # 提取边缘作为控制条件
    canny_image = cv2.Canny(np.array(image), 100, 200)
    
    # 生成缺陷图像
    result = pipe(
        prompt=prompt,
        image=image,
        mask_image=mask,
        control_image=canny_image,
        num_inference_steps=30,
        strength=0.8  # 控制生成强度
    ).images[0]
    
    return result

# 使用示例
defect_img = generate_defect(
    good_image_path="normal_part.jpg",
    mask_path="scratch_mask.png",  # 指定划痕位置
    prompt="a deep scratch with metallic luster on metal surface, industrial defect"
)
defect_img.save("generated_scratch.jpg")

import random
from pathlib import Path

class DefectGenerator:
    def __init__(self, pipe):
        self.pipe = pipe
    
    def batch_generate(self, good_dir, mask_dir, output_dir, prompt_template, num_per_class=100):
        """
        批量生成缺陷数据
        - good_dir: 良品图目录
        - mask_dir: 缺陷掩码目录（可准备多种掩码模板）
        - prompt_template: 提示词模板，如 "{defect_type} on {material} surface"
        """
        good_images = list(Path(good_dir).glob("*.jpg"))
        masks = list(Path(mask_dir).glob("*.png"))
        
        defect_types = ["scratch", "dent", "burr", "crack", "stain"]
        materials = ["metal", "plastic", "glass", "ceramic"]
        
        for i in range(num_per_class):
            # 随机组合
            good_img = random.choice(good_images)
            mask = random.choice(masks)
            defect_type = random.choice(defect_types)
            material = random.choice(materials)
            
            prompt = prompt_template.format(
                defect_type=defect_type, 
                material=material
            )
            
            result = self.pipe(
                prompt=prompt,
                image=load_image(good_img),
                mask_image=load_image(mask),
                num_inference_steps=30
            ).images[0]
            
            result.save(f"{output_dir}/{defect_type}_{i}.jpg")

import subprocess
from ultralytics import YOLO

class IndustrialDefectPipeline:
    """
    完整的工业缺陷检测流水线
    1. AIGC生成缺陷数据
    2. YOLO模型训练
    3. OpenClaw调度
    """
    
    def __init__(self, yolo_model_path="yolov8n.pt"):
        self.generator = DefectGenerator(pipe)
        self.yolo = YOLO(yolo_model_path)
    
    def run_full_pipeline(self, good_dir, mask_dir, output_dir, epochs=100):
        # Step 1: 生成缺陷数据
        print("Step 1: Generating defect images...")
        self.generator.batch_generate(
            good_dir, mask_dir, f"{output_dir}/generated", 
            prompt_template="{defect_type} on {material} surface, industrial defect"
        )
        
        # Step 2: 合并真实缺陷和生成缺陷
        print("Step 2: Merging datasets...")
        # 这里需要将真实缺陷图和生成缺陷图合并到一个目录
        # 同时生成对应的标注文件（可用前面的SAM自动标注）
        
        # Step 3: 训练YOLO模型
        print("Step 3: Training YOLO model...")
        results = self.yolo.train(
            data=f"{output_dir}/dataset.yaml",
            epochs=epochs,
            imgsz=640,
            batch=16,
            device=0
        )
        
        return results

用户：每周日凌晨2点，用 defect-generation 技能，
      先基于良品图生成5000张缺陷图，
      然后自动训练YOLO模型，
      最后用最新模型替换线上检测服务
OpenClaw：已创建定时任务，每周日2:00执行缺陷生成+模型训练+自动部署

# 安装依赖
pip install diffusers transformers accelerate xformers

# 下载ControlNet和Stable Diffusion模型
python -c "from diffusers import StableDiffusionControlNetPipeline; \
    pipe = StableDiffusionControlNetPipeline.from_pretrained('runwayml/stable-diffusion-v1-5')"