基于YOLO26的PCB表面缺陷检测系统（Python源码+数据集+Pyside6界面）

重复模块m (n次迭代):
┌─────────────────────────────────────────────────────────┐
│                                                         │
│  如果 attn=True:                                        │
│    Sequential(                                          │
│        Bottleneck(self.c, self.c),                     │  ←─ 先特征提取
│        PSABlock(self.c, attn_ratio=0.5, num_heads=...) │  ←─ 后注意力增强
│    )                                                    │
│                                                         │
│  否则如果 c3k=True:                                     │
│    C3k(self.c, self.c, 2)                             │  ←─ 同YOLOv11
│                                                         │
│  否则:                                                  │
│    Bottleneck(self.c, self.c)                         │  ←─ 同YOLOv11
│                                                         │
└─────────────────────────────────────────────────────────┘

class BboxLoss(nn.Module):
    """Criterion class for computing training losses for bounding boxes."""

    def __init__(self, reg_max: int = 16):
        """Initialize the BboxLoss module with regularization maximum and DFL settings."""
        super().__init__()
        self.dfl_loss = DFLoss(reg_max) if reg_max > 1 else None

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于YOLO26的PCB表面缺陷检测系统
功能：训练、推理（图像/视频/摄像头）、缺陷尺寸测量、结果可视化、报告生成
依赖：ultralytics, opencv-python, numpy, torch, matplotlib, scipy, scikit-image
"""

import os
import argparse
import time
import json
import math
from pathlib import Path
from typing import Union, List, Tuple, Dict, Optional

import cv2
import numpy as np
import torch
from ultralytics import YOLO
from scipy.spatial.distance import cdist
from skimage.measure import label, regionprops

# ======================= 配置参数 =======================
DEFAULT_CONF = 0.25          # 置信度阈值
DEFAULT_IOU = 0.45           # NMS IoU阈值
DEFAULT_IMGSZ = 640          # 推理图像尺寸
MM_PER_PIXEL = 0.05          # 像素当量：每个像素对应的实际尺寸（mm），根据相机标定调整
OUTPUT_DIR = "outputs"       # 输出结果保存目录

# PCB缺陷类别（示例，根据实际数据集修改）
CLASS_NAMES = {
    0: "open_circuit",    # 开路
    1: "short_circuit",   # 短路
    2: "burr",            # 毛刺
    3: "dent",            # 凹陷
    4: "scratch",         # 划痕
    5: "hole"             # 孔洞
}
# 类别对应的颜色（BGR）
CLASS_COLORS = {
    0: (0, 0, 255),       # 红色
    1: (0, 255, 0),       # 绿色
    2: (255, 0, 0),       # 蓝色
    3: (0, 255, 255),     # 黄色
    4: (255, 0, 255),     # 品红
    5: (255, 255, 0)      # 青色
}


# ======================= 缺陷量化工具 =======================
class DefectMeasurer:
    """缺陷尺寸测量工具（基于分割掩码）"""
    
    @staticmethod
    def pixel_to_mm(pixels: float, mm_per_pixel: float = MM_PER_PIXEL) -> float:
        """像素距离转实际毫米"""
        return pixels * mm_per_pixel
    
    @staticmethod
    def compute_mask_properties(mask: np.ndarray, mm_per_pixel: float = MM_PER_PIXEL) -> Dict:
        """
        计算分割掩码的几何属性
        Args:
            mask: 二值掩码 (H, W), dtype=uint8 or bool
            mm_per_pixel: 像素当量 (mm/pixel)
        Returns:
            dict: 包含面积、周长、等效直径、质心等
        """
        if mask.dtype != bool:
            mask = mask > 0
        props = regionprops(mask.astype(np.uint8))[0] if np.any(mask) else None
        if props is None:
            return {}
        # 像素单位属性
        area_px = props.area
        perimeter_px = props.perimeter
        equiv_diameter_px = np.sqrt(4 * area_px / np.pi)
        centroid_px = props.centroid  # (row, col)
        # 转换为毫米单位
        area_mm = area_px * mm_per_pixel ** 2
        perimeter_mm = perimeter_px * mm_per_pixel
        equiv_diameter_mm = equiv_diameter_px * mm_per_pixel
        return {
            "area_px": area_px,
            "area_mm2": area_mm,
            "perimeter_px": perimeter_px,
            "perimeter_mm": perimeter_mm,
            "equiv_diameter_px": equiv_diameter_px,
            "equiv_diameter_mm": equiv_diameter_mm,
            "centroid_px": centroid_px,
            "bbox_px": props.bbox,  # (min_row, min_col, max_row, max_col)
        }


# ======================= 可视化与报告 =======================
class Visualizer:
    """检测结果可视化工具"""
    
    @staticmethod
    def draw_boxes(image: np.ndarray, boxes: np.ndarray, confs: np.ndarray,
                   cls_ids: np.ndarray, masks: Optional[np.ndarray] = None) -> np.ndarray:
        """
        在图像上绘制检测框、掩码、置信度、类别标签
        Args:
            image: 原始图像 (BGR)
            boxes: (N, 4) 边界框 [x1,y1,x2,y2] 整数坐标
            confs: (N,) 置信度
            cls_ids: (N,) 类别ID
            masks: (N, H, W) 二值掩码（可选）
        Returns:
            带标注的图像
        """
        result = image.copy()
        # 先画掩码（半透明填充）
        if masks is not None:
            overlay = result.copy()
            for i, mask in enumerate(masks):
                if mask is None or mask.sum() == 0:
                    continue
                color = CLASS_COLORS.get(cls_ids[i], (128, 128, 128))
                # 将mask转为3通道
                mask_3c = np.stack([mask]*3, axis=-1).astype(np.uint8) * 255
                colored_mask = np.zeros_like(overlay)
                colored_mask[mask] = color
                overlay = cv2.addWeighted(overlay, 0.6, colored_mask, 0.4, 0)
            result = overlay
        
        # 画边框和标签
        for i, (box, conf, cls_id) in enumerate(zip(boxes, confs, cls_ids)):
            x1, y1, x2, y2 = box
            color = CLASS_COLORS.get(cls_id, (128, 128, 128))
            cv2.rectangle(result, (x1, y1), (x2, y2), color, 2)
            label = f"{CLASS_NAMES.get(cls_id, 'unknown')}: {conf:.2f}"
            (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.rectangle(result, (x1, y1 - th - 5), (x1 + tw, y1), color, -1)
            cv2.putText(result, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
        return result
    
    @staticmethod
    def save_report(detections: List[Dict], output_path: str, image_name: str):
        """
        保存检测报告为JSON文件
        Args:
            detections: 检测结果列表，每个元素包含 'cls', 'conf', 'bbox', 'mask_properties' 等
            output_path: 报告保存路径
            image_name: 图像文件名
        """
        report = {
            "image": image_name,
            "detections": detections,
            "summary": {
                "total_defects": len(detections),
                "defect_types": {}
            }
        }
        for d in detections:
            cls_name = d['class']
            report['summary']['defect_types'][cls_name] = report['summary']['defect_types'].get(cls_name, 0) + 1
        with open(output_path, 'w') as f:
            json.dump(report, f, indent=2)
        print(f"报告已保存至 {output_path}")


# ======================= 检测器主类 =======================
class PCBDefectDetector:
    """PCB缺陷检测器，基于YOLO26"""
    
    def __init__(self, model_path: str, conf_thres: float = DEFAULT_CONF,
                 iou_thres: float = DEFAULT_IOU, device: str = None):
        """
        初始化检测器
        Args:
            model_path: 模型权重路径 (.pt)
            conf_thres: 置信度阈值
            iou_thres: NMS IoU阈值
            device: 推理设备 ('cpu', 'cuda:0' 等)
        """
        self.model = YOLO(model_path)
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        if device is None:
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.device = device
        self.measurer = DefectMeasurer()
        self.visualizer = Visualizer()
    
    def predict(self, source: Union[str, np.ndarray], save_dir: str = None,
                visualize: bool = True, measure: bool = True) -> List[Dict]:
        """
        对图像或视频源进行检测
        Args:
            source: 图像路径、视频路径、摄像头索引(如0) 或 numpy图像 (BGR)
            save_dir: 输出结果保存目录，None表示不保存
            visualize: 是否显示检测结果窗口
            measure: 是否进行缺陷尺寸测量
        Returns:
            检测结果列表
        """
        # 处理源类型
        if isinstance(source, str):
            if source.isdigit():
                source = int(source)  # 摄像头
            elif source.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
                source = source  # 视频文件
            else:
                source = source  # 图像文件
        # 调用模型预测
        results = self.model.predict(
            source=source,
            conf=self.conf_thres,
            iou=self.iou_thres,
            imgsz=DEFAULT_IMGSZ,
            device=self.device,
            stream=False if isinstance(source, np.ndarray) else True,
            verbose=False
        )
        
        all_detections = []
        # 处理每帧/每图
        for result in results:
            img = result.orig_img  # BGR numpy
            boxes = result.boxes
            masks = result.masks  # 分割掩码 (如果有)
            if boxes is None:
                detections = []
                annotated = img
            else:
                # 提取数据
                xyxy = boxes.xyxy.cpu().numpy().astype(int)  # Nx4
                confs = boxes.conf.cpu().numpy()  # N
                cls_ids = boxes.cls.cpu().numpy().astype(int)  # N
                # 提取掩码（如果有）
                mask_array = None
                if masks is not None:
                    # masks.data shape: (N, H, W) 或 (N, 160, 160)
                    mask_array = masks.data.cpu().numpy()
                    # 如果掩码是缩小的，需要上采样到原始图像尺寸
                    if mask_array.shape[1] != img.shape[0] or mask_array.shape[2] != img.shape[1]:
                        mask_array = np.array([cv2.resize(m, (img.shape[1], img.shape[0])) for m in mask_array])
                # 构建检测结果列表并测量尺寸
                detections = []
                for i, (box, conf, cls_id) in enumerate(zip(xyxy, confs, cls_ids)):
                    det = {
                        "class": CLASS_NAMES.get(cls_id, "unknown"),
                        "class_id": int(cls_id),
                        "confidence": float(conf),
                        "bbox": box.tolist(),
                    }
                    if measure and mask_array is not None and i < len(mask_array):
                        mask = (mask_array[i] > 0).astype(np.uint8)
                        props = self.measurer.compute_mask_properties(mask, MM_PER_PIXEL)
                        det["mask_properties"] = props
                    detections.append(det)
                # 可视化
                if visualize:
                    # 提取掩码用于绘制
                    draw_masks = None
                    if mask_array is not None:
                        draw_masks = (mask_array > 0).astype(bool)
                    annotated = self.visualizer.draw_boxes(img, xyxy, confs, cls_ids, draw_masks)
                else:
                    annotated = img
                all_detections.append(detections)
            # 显示窗口
            if visualize:
                cv2.imshow("PCB Defect Detection", annotated)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            # 保存输出
            if save_dir:
                os.makedirs(save_dir, exist_ok=True)
                # 如果source是文件名，提取名称
                if isinstance(source, str):
                    name = os.path.basename(source).split('.')[0]
                else:
                    name = f"frame_{int(time.time())}"
                img_path = os.path.join(save_dir, f"{name}_detected.jpg")
                cv2.imwrite(img_path, annotated)
                print(f"检测结果保存至 {img_path}")
                # 保存测量报告
                if detections:
                    report_path = os.path.join(save_dir, f"{name}_report.json")
                    self.visualizer.save_report(detections, report_path, name + '.jpg')
        if visualize:
            cv2.destroyAllWindows()
        return all_detections


# ======================= 训练函数 =======================
def train_model(config_yaml: str, model_name: str = 'yolo26s.pt', epochs: int = 200,
                imgsz: int = DEFAULT_IMGSZ, batch: int = 16, device: str = None):
    """
    训练YOLO26模型
    Args:
        config_yaml: data.yaml 路径
        model_name: 预训练模型名称，如 'yolo26n.pt', 'yolo26s.pt'
        epochs: 训练轮数
        imgsz: 输入图像尺寸
        batch: 批量大小
        device: 设备
    """
    model = YOLO(model_name)
    results = model.train(
        data=config_yaml,
        epochs=epochs,
        imgsz=imgsz,
        batch=batch,
        device=device,
        workers=8,
        lr0=0.01,
        weight_decay=0.0005,
        warmup_epochs=3,
        cos_lr=True,
        augment=True,
        seed=42,
        verbose=True
    )
    print(f"训练完成，最佳模型保存在: {model.trainer.best}")
    return model


# ======================= 主程序入口 =======================
def main():
    parser = argparse.ArgumentParser(description="基于YOLO26的PCB表面缺陷检测系统")
    subparsers = parser.add_subparsers(dest='mode', help='运行模式')
    
    # 训练模式
    train_parser = subparsers.add_parser('train', help='训练模型')
    train_parser.add_argument('--data', type=str, required=True, help='数据集配置文件 data.yaml')
    train_parser.add_argument('--model', type=str, default='yolo26s.pt', help='预训练模型')
    train_parser.add_argument('--epochs', type=int, default=200, help='训练轮数')
    train_parser.add_argument('--batch', type=int, default=16, help='批量大小')
    train_parser.add_argument('--device', type=str, default=None, help='设备(cpu/cuda)')
    
    # 检测模式
    detect_parser = subparsers.add_parser('detect', help='检测缺陷')
    detect_parser.add_argument('--source', type=str, required=True, help='输入源：图像路径/视频路径/摄像头索引(0,1...)')
    detect_parser.add_argument('--weights', type=str, required=True, help='模型权重路径')
    detect_parser.add_argument('--conf', type=float, default=DEFAULT_CONF, help='置信度阈值')
    detect_parser.add_argument('--iou', type=float, default=DEFAULT_IOU, help='NMS IoU阈值')
    detect_parser.add_argument('--save-dir', type=str, default=OUTPUT_DIR, help='结果保存目录')
    detect_parser.add_argument('--no-show', action='store_true', help='不显示检测窗口')
    detect_parser.add_argument('--no-measure', action='store_true', help='不进行尺寸测量')
    
    args = parser.parse_args()
    
    if args.mode == 'train':
        train_model(args.data, args.model, args.epochs, DEFAULT_IMGSZ, args.batch, args.device)
    elif args.mode == 'detect':
        detector = PCBDefectDetector(
            model_path=args.weights,
            conf_thres=args.conf,
            iou_thres=args.iou
        )
        # 处理源类型（数字字符串转整数表示摄像头）
        source = args.source
        if source.isdigit():
            source = int(source)
        detector.predict(
            source=source,
            save_dir=args.save_dir,
            visualize=not args.no_show,
            measure=not args.no_measure
        )
    else:
        parser.print_help()


if __name__ == '__main__':
    main()

nc: 9
names:
- SH
- SP
- SC
- OP
- MB
- HB
- CS
- CFO
- BMFO

path: D:/YOLO26/data/PCB/
train: images/train
val: images/val
nc: 9
names:
- SH
- SP
- SC
- OP
- MB
- HB
- CS
- CFO
- BMFO

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/26/yolo26.yaml')
    model.train(data='data/PCB.yaml',
                cache=False,
                imgsz=640,
                epochs=200,
                batch=8,
                close_mosaic=10,
                device='0',
                optimizer='SGD', # using SGD
                project='runs/train',
                name='exp',
                )

YOLO26 summary (fused): 122 layers, 2,376,591 parameters, 0 gradients, 5.2 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% ━━━━━━━━━━━━ 65/65 3.6it/s 17.9s
                   all       2051       4092      0.826      0.771      0.842      0.509
                    SH        126        169      0.825      0.858      0.909      0.592
                    SP        430        929      0.821      0.736      0.819      0.381
                    SC        245        285      0.808      0.747      0.824      0.507
                    OP        274        338      0.837      0.817      0.896      0.545
                    MB        391        546       0.83       0.69      0.804      0.386
                    HB        271        608      0.935      0.944      0.983      0.822
                    CS        279        448      0.766      0.641      0.736      0.461
                   CFO        309        423      0.771       0.65      0.718      0.415
                  BMFO        305        346       0.84      0.855      0.889      0.473

pip install --upgrade pip
pip install pyside6 -i https://mirror.baidu.com/pypi/simple