工业质检的“四座大山”：样本少、不平衡、拍不全、导入慢——YOLO+无监督组合拳，漏检率直降96%

"""
工业缺陷检测双引擎方案
作者：AI小怪兽
技术栈：YOLOv8 (有监督) + PatchCore (无监督) + 模拟光学成像数据增强
功能：1. 仅用良品图像训练无监督模型（快速冷启动）
      2. 利用模拟光学成像增强少量缺陷样本，训练YOLO模型
      3. 双引擎协同推理，输出最终检测结果
依赖：pip install ultralytics albumentations anomalib opencv-python torch
"""

import cv2
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
import torch
from ultralytics import YOLO
from anomalib.models import Patchcore
from anomalib.data import Folder
from anomalib.engine import Engine
from pathlib import Path

# ======================== 配置参数 ========================
GOOD_IMG_DIR = "./datasets/good"           # 良品图像目录
DEFECT_DIR = "./datasets/defect"           # 缺陷图像目录（按类别分文件夹，如 defect/scratch/）
YOLO_DATA_YAML = "dataset.yaml"            # YOLO数据集配置文件
PATCHCORE_CKPT = "patchcore_model.ckpt"    # 无监督模型保存路径
YOLO_MODEL_PATH = "yolov8n_custom.pt"      # 训练后的YOLO模型路径
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
IMG_SIZE = 640

# ======================== 1. 模拟光学成像数据增强 ========================
def get_optical_augmentation():
    """返回模拟光学成像的增强流水线（仅训练时使用）"""
    return A.Compose([
        # 光照变化
        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
        A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=20, val_shift_limit=10, p=0.5),
        # 镜头畸变（透视变换）
        A.Perspective(scale=(0.05, 0.1), p=0.3),
        # 材质反光（模拟高光）
        A.RandomSunFlare(flare_roi=(0, 0, 1, 0.5), src_radius=100, p=0.3),
        # 成像噪声
        A.GaussNoise(var_limit=(10.0, 50.0), p=0.4),
        A.MotionBlur(blur_limit=5, p=0.3),
        # 保证输出尺寸一致
        A.Resize(IMG_SIZE, IMG_SIZE),
    ])

# ======================== 2. 无监督异常检测模型训练（PatchCore） ========================
def train_unsupervised():
    """仅使用良品图像训练无监督模型，实现零缺陷样本冷启动"""
    print(">> 开始训练无监督模型（PatchCore）...")
    # 数据模块：只读取良品图像
    datamodule = Folder(
        name="industrial_good",
        root="./datasets",
        normal_dir="good",
        task="segmentation",           # 输出像素级异常图
        image_size=(IMG_SIZE, IMG_SIZE),
    )
    model = Patchcore(
        backbone="resnet50",
        pre_trained=True,
        coreset_sampling_ratio=0.1,
    )
    engine = Engine(device=DEVICE)
    engine.fit(datamodule=datamodule, model=model)
    engine.save_checkpoint(PATCHCORE_CKPT)
    print(f"无监督模型已保存至 {PATCHCORE_CKPT}")

# ======================== 3. YOLOv8 训练（自定义数据集，嵌入模拟光学成像增强） ========================
# 3.1 自定义YOLO数据集类，覆盖数据增强方法
from ultralytics.yolo.data.dataset import YOLODataset

class OpticalAugYOLODataset(YOLODataset):
    """在YOLO数据集基础上增加模拟光学成像增强"""
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.optical_aug = get_optical_augmentation()

    def load_image(self, i, rect_mode=True):
        """加载图像并应用模拟光学成像增强（仅训练阶段）"""
        img, (h0, w0), (h, w) = super().load_image(i, rect_mode)
        if self.augment:
            # 应用我们的光学模拟增强
            augmented = self.optical_aug(image=img)
            img = augmented["image"]
        return img, (h0, w0), (h, w)

def train_yolo():
    """使用少量缺陷样本 + 模拟光学成像增强训练YOLOv8"""
    print(">> 开始训练YOLOv8（有监督）...")
    # 替换默认的数据集类
    from ultralytics.yolo.data import build_dataset
    build_dataset.YOLODataset = OpticalAugYOLODataset

    model = YOLO("yolov8n.pt")
    model.train(
        data=YOLO_DATA_YAML,
        epochs=100,
        imgsz=IMG_SIZE,
        batch=16,
        device=DEVICE,
        augment=True,                  # 启用增强（会调用我们的custom dataset）
        patience=10,
        save=True,
        project="runs/train",
        name="yolo_optical_aug",
    )
    # 保存最终模型
    model.save(YOLO_MODEL_PATH)
    print(f"YOLO模型已保存至 {YOLO_MODEL_PATH}")

# ======================== 4. 双引擎协同推理 ========================
class DualEngineInference:
    """加载两个模型，协同进行缺陷检测"""
    def __init__(self, yolo_path, patchcore_ckpt):
        self.device = DEVICE
        self.yolo = YOLO(yolo_path)
        # 加载无监督模型
        self.anomaly_model = Patchcore.load_from_checkpoint(patchcore_ckpt)
        self.anomaly_model.eval()
        self.anomaly_model.to(self.device)

    def predict(self, image_path):
        """对单张图像进行检测，返回 (缺陷类别, 置信度, 异常区域mask)"""
        img = cv2.imread(image_path)
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        orig_h, orig_w = img.shape[:2]

        # 步骤1：YOLO检测已知缺陷
        yolo_results = self.yolo(img_rgb, imgsz=IMG_SIZE, conf=0.25, device=self.device)[0]
        yolo_defects = []   # 存储 (bbox, cls_name, conf)
        for box in yolo_results.boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            cls_name = self.yolo.names[cls_id]
            bbox = box.xyxy[0].cpu().numpy().astype(int)
            yolo_defects.append((bbox, cls_name, conf))

        # 步骤2：无监督异常检测（像素级）
        # 需要将图像调整为模型输入尺寸
        input_tensor = self._preprocess_for_anomaly(img_rgb)
        with torch.no_grad():
            predictions = self.anomaly_model(input_tensor)
        anomaly_map = predictions["anomaly_map"].squeeze().cpu().numpy()  # (H, W)
        # 阈值判断是否有未知异常（可根据实际调整）
        unknown_anomaly = np.any(anomaly_map > 0.5)   # 0.5为经验阈值
        if unknown_anomaly:
            # 生成二值mask用于定位
            mask = (anomaly_map > 0.5).astype(np.uint8) * 255
            # 可以再根据mask寻找轮廓等，这里简化为整体标记
            unknown_defect = ("unknown", 1.0, mask)
        else:
            unknown_defect = None

        return {
            "yolo_defects": yolo_defects,
            "unknown_defect": unknown_defect,
            "anomaly_map": anomaly_map
        }

    def _preprocess_for_anomaly(self, img_rgb):
        """将图像转换为模型输入格式（归一化、tensor等）"""
        # 假设无监督模型输入尺寸与训练时一致（这里是IMG_SIZE）
        img = cv2.resize(img_rgb, (IMG_SIZE, IMG_SIZE))
        img = img.astype(np.float32) / 255.0
        img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
        return img.to(self.device)

# ======================== 5. 主流程示例 ========================
if __name__ == "__main__":
    # 第一步：训练无监督模型（仅需良品图像，若已存在可跳过）
    # train_unsupervised()

    # 第二步：训练YOLO模型（需要少量缺陷样本+标注）
    # train_yolo()

    # 第三步：加载双引擎进行推理演示
    detector = DualEngineInference(YOLO_MODEL_PATH, PATCHCORE_CKPT)
    test_image = "./test_images/sample.jpg"
    result = detector.predict(test_image)

    print("YOLO检测到的已知缺陷：")
    for bbox, cls, conf in result["yolo_defects"]:
        print(f"  - {cls}: confidence={conf:.2f}, bbox={bbox}")

    if result["unknown_defect"]:
        cls, conf, mask = result["unknown_defect"]
        print(f"无监督模型发现未知缺陷: {cls}, confidence={conf:.2f}, mask区域已定位")
    else:
        print("无监督模型未发现异常")

# ======================== 说明 ========================
"""
关键点：
1. 模拟光学成像增强：通过albumentations模拟光照、畸变、反光、噪声，使YOLO模型对“拍不全”问题鲁棒。
2. 无监督模型仅用良品图像训练，快速上线，自动积累未知缺陷样本。
3. 双引擎协同：YOLO负责已知缺陷分类，无监督负责兜底未知缺陷，形成闭环。
4. 通过自定义YOLO数据集类，完美嵌入光学增强，无需修改ultralytics内核。
"""