基于YOLO目标检测模型的大规模生产电子元件自动化缺陷检测

# =============================================================================
# Automated Defect Detection for Mass-Produced Electronic Components
# Based on YOLO Object Detection Models - 伪代码实现
# =============================================================================

import cv2
import numpy as np
import torch
import yolov5  # 或其他YOLO版本

class ElectronicComponentDefectDetector:
    def __init__(self, model_path, conf_threshold=0.5, iou_threshold=0.4):
        """
        初始化缺陷检测器
        Args:
            model_path: 预训练的YOLO模型路径
            conf_threshold: 置信度阈值
            iou_threshold: IoU阈值用于NMS
        """
        self.model = yolov5.load(model_path)
        self.model.conf = conf_threshold
        self.model.iou = iou_threshold
        self.defect_classes = ['short_circuit', 'open_circuit', 'solder_void', 
                              'component_misalignment', 'crack', 'contamination']
        
    def preprocess_image(self, image_path):
        """
        图像预处理
        """
        # 读取图像
        image = cv2.imread(image_path)
        # 转换为RGB
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 图像增强 (可选)
        enhanced_image = self.image_enhancement(image_rgb)
        return enhanced_image
    
    def image_enhancement(self, image):
        """
        图像增强以提高检测精度
        """
        # 直方图均衡化
        lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
        lab[:,:,0] = cv2.equalizeHist(lab[:,:,0])
        enhanced = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
        
        # 高斯模糊去噪
        denoised = cv2.GaussianBlur(enhanced, (3,3), 0)
        
        return denoised
    
    def detect_defects(self, image_path):
        """
        执行缺陷检测
        """
        # 1. 图像预处理
        processed_image = self.preprocess_image(image_path)
        
        # 2. YOLO模型推理
        results = self.model(processed_image)
        
        # 3. 后处理
        detections = results.pandas().xyxy[0]  # 获取检测结果
        
        # 4. 分析检测结果
        defect_analysis = self.analyze_defects(detections)
        
        return detections, defect_analysis
    
    def analyze_defects(self, detections):
        """
        分析缺陷检测结果
        """
        analysis = {
            'total_defects': len(detections),
            'defect_types': {},
            'severity_level': 'PASS',
            'defect_locations': []
        }
        
        for _, detection in detections.iterrows():
            defect_class = detection['class']
            confidence = detection['confidence']
            
            # 统计各类缺陷数量
            if defect_class in analysis['defect_types']:
                analysis['defect_types'][defect_class] += 1
            else:
                analysis['defect_types'][defect_class] = 1
            
            # 记录缺陷位置
            bbox = {
                'xmin': detection['xmin'],
                'ymin': detection['ymin'],
                'xmax': detection['xmax'],
                'ymax': detection['ymax']
            }
            analysis['defect_locations'].append({
                'class': defect_class,
                'confidence': confidence,
                'bbox': bbox
            })
        
        # 根据缺陷数量和类型确定严重程度
        if analysis['total_defects'] > 3:
            analysis['severity_level'] = 'CRITICAL'
        elif analysis['total_defects'] > 0:
            analysis['severity_level'] = 'WARNING'
        
        return analysis
    
    def visualize_results(self, image_path, detections, save_path=None):
        """
        可视化检测结果
        """
        image = cv2.imread(image_path)
        
        for _, detection in detections.iterrows():
            # 绘制边界框
            xmin, ymin, xmax, ymax = int(detection['xmin']), int(detection['ymin']), \
                                   int(detection['xmax']), int(detection['ymax'])
            
            # 根据缺陷类型选择颜色
            color = self.get_defect_color(detection['class'])
            
            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)
            
            # 添加标签
            label = f"{detection['name']}: {detection['confidence']:.2f}"
            cv2.putText(image, label, (xmin, ymin-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        
        if save_path:
            cv2.imwrite(save_path, image)
        
        return image
    
    def get_defect_color(self, defect_class):
        """
        根据缺陷类型返回对应的颜色
        """
        color_map = {
            'short_circuit': (0, 0, 255),      # 红色 - 严重缺陷
            'open_circuit': (0, 165, 255),     # 橙色 - 严重缺陷
            'crack': (0, 255, 255),           # 黄色 - 中等缺陷
            'solder_void': (255, 255, 0),     # 青色 - 轻微缺陷
            'component_misalignment': (255, 0, 0), # 蓝色 - 轻微缺陷
            'contamination': (128, 0, 128)     # 紫色 - 轻微缺陷
        }
        return color_map.get(defect_class, (255, 255, 255))

# =============================================================================
# 训练过程的伪代码
# =============================================================================

def train_defect_detection_model():
    """
    训练YOLO缺陷检测模型
    """
    # 1. 数据准备
    dataset_config = {
        'train': 'dataset/train/images',
        'val': 'dataset/val/images',
        'test': 'dataset/test/images',
        'nc': 6,  # 缺陷类别数量
        'names': ['short_circuit', 'open_circuit', 'solder_void', 
                 'component_misalignment', 'crack', 'contamination']
    }
    
    # 2. 模型配置
    model_config = {
        'model_type': 'yolov5s',  # 可选择yolov5n, yolov5s, yolov5m, yolov5l, yolov5x
        'img_size': 640,
        'batch_size': 16,
        'epochs': 100,
        'learning_rate': 0.01,
        'optimizer': 'SGD',
        'momentum': 0.937,
        'weight_decay': 0.0005
    }
    
    # 3. 数据增强配置
    augmentation_config = {
        'hsv_h': 0.015,  # 色调增强
        'hsv_s': 0.7,    # 饱和度增强
        'hsv_v': 0.4,    # 明度增强
        'rotate': 10,    # 旋转角度
        'translate': 0.1, # 平移
        'scale': 0.5,    # 缩放
        'shear': 0.0,    # 剪切
        'flipud': 0.0,   # 上下翻转
        'fliplr': 0.5,   # 左右翻转
        'mosaic': 1.0,   # Mosaic数据增强
        'mixup': 0.0     # Mixup数据增强
    }
    
    # 4. 训练过程
    for epoch in range(model_config['epochs']):
        # 训练阶段
        model.train()
        for batch_idx, (images, targets) in enumerate(train_loader):
            # 前向传播
            predictions = model(images)
            
            # 计算损失
            loss = compute_loss(predictions, targets)
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # 记录训练指标
            log_training_metrics(loss, batch_idx)
        
        # 验证阶段
        model.eval()
        with torch.no_grad():
            for images, targets in val_loader:
                predictions = model(images)
                # 计算验证指标
                calculate_validation_metrics(predictions, targets)
        
        # 保存最佳模型
        if is_best_model(current_metrics):
            save_checkpoint(model, epoch, current_metrics)

# =============================================================================
# 主程序
# =============================================================================

def main():
    """
    主执行函数
    """
    # 初始化缺陷检测器
    detector = ElectronicComponentDefectDetector(
        model_path='best_defect_detection_model.pt',
        conf_threshold=0.6,
        iou_threshold=0.45
    )
    
    # 批量处理电子组件图像
    component_images = [
        'production_batch_001/component_001.jpg',
        'production_batch_001/component_002.jpg',
        # ... 更多图像
    ]
    
    results = []
    for image_path in component_images:
        # 执行缺陷检测
        detections, analysis = detector.detect_defects(image_path)
        
        # 可视化结果
        result_image = detector.visualize_results(
            image_path, 
            detections, 
            save_path=f'results/{os.path.basename(image_path)}'
        )
        
        # 保存检测结果
        results.append({
            'image_path': image_path,
            'analysis': analysis,
            'defect_count': analysis['total_defects'],
            'severity': analysis['severity_level']
        })
        
        # 打印检测结果
        print(f"图像: {image_path}")
        print(f"缺陷数量: {analysis['total_defects']}")
        print(f"严重程度: {analysis['severity_level']}")
        print("---")
    
    # 生成检测报告
    generate_inspection_report(results)

def generate_inspection_report(results):
    """
    生成质量检测报告
    """
    total_components = len(results)
    defective_components = sum(1 for r in results if r['defect_count'] > 0)
    yield_rate = (total_components - defective_components) / total_components * 100
    
    print(f"\n=== 质量检测报告 ===")
    print(f"总检测数量: {total_components}")
    print(f"缺陷产品数量: {defective_components}")
    print(f"良率: {yield_rate:.2f}%")
    
    # 统计各类缺陷
    defect_statistics = {}
    for result in results:
        for defect_type, count in result['analysis']['defect_types'].items():
            if defect_type in defect_statistics:
                defect_statistics[defect_type] += count
            else:
                defect_statistics[defect_type] = count
    
    print("\n缺陷类型统计:")
    for defect_type, count in defect_statistics.items():
        print(f"  {defect_type}: {count}")

if __name__ == "__main__":
    main()