基于YOLO26的透明塑料瓶检测系统（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

# ======================================================================
# 基于YOLO26的透明塑料瓶检测系统 - 伪代码
# 数据集：Plastic Bottle Component & Defect Inspection Dataset
# 类别：bottle, cap, label, liquid
# 功能：多组件检测、缺陷筛查、几何质量验证
# ======================================================================

import os
import cv2
import yaml
import torch
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from ultralytics import YOLO
from typing import List, Dict, Tuple

# ========================= 1. 配置参数 =========================
class Config:
    """系统配置参数"""
    # 数据集
    DATA_YAML = "bottle_detect.yaml"         # YOLO检测配置文件
    CLASS_NAMES = ["bottle", "cap", "label", "liquid"]
    NUM_CLASSES = 4

    # 模型
    MODEL_NAME = "yolo26n.pt"                # 预训练权重 (n/s/m/l/x)
    IMG_SIZE = 640                           # 输入尺寸
    BATCH_SIZE = 16
    EPOCHS = 100
    DEVICE = 0                               # GPU编号，-1为CPU

    # 路径
    PROJECT_DIR = "./bottle_detection"
    WEIGHT_DIR = f"{PROJECT_DIR}/weights"
    RESULT_DIR = f"{PROJECT_DIR}/results"

# ========================= 2. 数据集准备 =========================
class BottleDataset:
    """透明塑料瓶数据集管理"""

    def __init__(self, data_yaml: str):
        self.data_yaml = data_yaml
        self._load_config()

    def _load_config(self):
        with open(self.data_yaml, 'r') as f:
            self.cfg = yaml.safe_load(f)
        self.train_path = self.cfg['train']
        self.val_path = self.cfg['val']
        self.nc = self.cfg['nc']
        self.names = self.cfg['names']
        print(f"[数据集] 加载成功: {self.nc}类 -> {self.names}")

    def analyze_defect_scenarios(self):
        """分析缺陷场景分布（可选）"""
        # 可根据文件路径或额外标签统计缺失瓶盖/标签等场景
        pass

# ========================= 3. 模型训练与推理 =========================
class BottleDetector:
    """透明塑料瓶多组件检测器"""

    def __init__(self, model_name: str = Config.MODEL_NAME):
        self.model_name = model_name
        self.model = None
        self._load_model()

    def _load_model(self):
        """加载YOLO26预训练模型"""
        self.model = YOLO(self.model_name)
        print(f"[模型] 加载预训练权重: {self.model_name}")

    def train(self, data_yaml: str, config: Config):
        """训练检测模型"""
        print("\n[训练] 开始训练透明塑料瓶检测模型...")

        train_args = {
            'data': data_yaml,
            'epochs': config.EPOCHS,
            'imgsz': config.IMG_SIZE,
            'batch': config.BATCH_SIZE,
            'device': config.DEVICE,
            'project': config.PROJECT_DIR,
            'name': 'yolo26_bottle',
            'exist_ok': True,
            'patience': 30,                # 早停
            'lr0': 0.01,
            'lrf': 0.01,
            'momentum': 0.937,
            'weight_decay': 0.0005,
            'warmup_epochs': 3,
            'box': 7.5,
            'cls': 0.5,
            'dfl': 1.5,
            'val': True,
            'plots': True,
            'save': True,
            'seed': 0,
            # 针对透明物体及工业场景的增强
            'mosaic': 0.5,
            'mixup': 0.0,
            'hsv_h': 0.015,                # 色调变化（适应不同光照）
            'hsv_s': 0.7,                  # 饱和度变化
            'hsv_v': 0.4,                  # 明度变化（应对镜面反射）
            'translate': 0.1,
            'scale': 0.5,
            'fliplr': 0.5,
            'flipud': 0.0,
            'degrees': 0.0,
        }

        results = self.model.train(**train_args)
        print("[训练] 训练完成!")
        return results

    def evaluate(self, data_yaml: str):
        """在验证集上评估模型"""
        print("\n[评估] 在验证集上评估模型...")

        metrics = self.model.val(
            data=data_yaml,
            split='val',
            conf=0.001,
            iou=0.6,
            device=Config.DEVICE,
            plots=True,
        )

        results = {
            'mAP50': metrics.box.map50,
            'mAP50-95': metrics.box.map,
            'precision': metrics.box.mp,
            'recall': metrics.box.mr,
            'F1': 2 * metrics.box.mp * metrics.box.mr / (metrics.box.mp + metrics.box.mr + 1e-6),
        }

        print(f"\n[评估结果]")
        print(f"  mAP@0.5: {results['mAP50']:.4f}")
        print(f"  mAP@0.5:0.95: {results['mAP50-95']:.4f}")
        print(f"  Precision: {results['precision']:.4f}")
        print(f"  Recall: {results['recall']:.4f}")
        print(f"  F1-Score: {results['F1']:.4f}")

        return results

    def predict(self, image_path: str, conf_thresh: float = 0.25):
        """单张图像推理"""
        results = self.model(image_path, conf=conf_thresh, iou=0.45)[0]
        detections = self._parse_results(results, image_path)
        return detections

    def predict_batch(self, image_dir: str, output_dir: str, conf_thresh: float = 0.25):
        """批量推理并保存可视化结果"""
        os.makedirs(output_dir, exist_ok=True)
        image_paths = list(Path(image_dir).glob("*.jpg")) + list(Path(image_dir).glob("*.png"))

        print(f"[批量推理] 共 {len(image_paths)} 张图像...")
        all_detections = []

        for img_path in image_paths:
            results = self.model(str(img_path), conf=conf_thresh)[0]
            annotated = results.plot()
            out_path = os.path.join(output_dir, img_path.name)
            cv2.imwrite(out_path, annotated)

            det = self._parse_results(results, str(img_path))
            all_detections.append(det)

        self._generate_report(all_detections, output_dir)
        return all_detections

    def _parse_results(self, results, image_path: str):
        """解析YOLO检测结果"""
        detections = []
        if results.boxes is not None:
            for box in results.boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                conf = float(box.conf[0])
                cls_id = int(box.cls[0])
                cls_name = results.names[cls_id]
                detections.append({
                    'class_id': cls_id,
                    'class_name': cls_name,
                    'confidence': conf,
                    'bbox': [x1, y1, x2, y2],
                })
        return {
            'image': Path(image_path).name,
            'num_detections': len(detections),
            'detections': detections,
        }

    def _generate_report(self, detections: List[Dict], output_dir: str):
        """生成检测统计报告，并计算几何质量指标"""
        report_path = os.path.join(output_dir, "detection_report.txt")
        total = sum(d['num_detections'] for d in detections)
        class_counts = {name: 0 for name in Config.CLASS_NAMES}

        # 收集每个图像的组件信息用于质量评估
        image_components = {}
        for d in detections:
            img_name = d['image']
            image_components[img_name] = {
                'bottle': None, 'cap': None, 'label': None, 'liquid': None
            }
            for det in d['detections']:
                cls_name = det['class_name']
                class_counts[cls_name] += 1
                if cls_name in image_components[img_name]:
                    image_components[img_name][cls_name] = det

        # 简单质量评估：缺失组件检测、液位填充率等
        anomalies = []
        for img_name, comps in image_components.items():
            bottle = comps['bottle']
            cap = comps['cap']
            label = comps['label']
            liquid = comps['liquid']

            # 缺失组件
            if bottle and not cap:
                anomalies.append(f"{img_name}: 缺失瓶盖")
            if bottle and not label:
                anomalies.append(f"{img_name}: 缺失标签")

            # 液位填充率估计（基于高度比）
            if bottle and liquid:
                bottle_h = bottle['bbox'][3] - bottle['bbox'][1]
                liquid_h = liquid['bbox'][3] - liquid['bbox'][1]
                if bottle_h > 0:
                    ratio = liquid_h / bottle_h
                    if ratio < 0.70:
                        anomalies.append(f"{img_name}: 液位不足 (比例 {ratio:.2f})")
                    elif ratio > 0.80:
                        anomalies.append(f"{img_name}: 液位过满 (比例 {ratio:.2f})")

        # 写入报告
        with open(report_path, 'w') as f:
            f.write("=" * 60 + "\n")
            f.write("透明塑料瓶检测报告\n")
            f.write("=" * 60 + "\n")
            f.write(f"处理图像数: {len(detections)}\n")
            f.write(f"检测组件总数: {total}\n\n")
            f.write("各组件统计:\n")
            for name, cnt in class_counts.items():
                if cnt > 0:
                    f.write(f"  {name}: {cnt}\n")
            f.write("\n检测到的异常:\n")
            for a in anomalies:
                f.write(f"  - {a}\n")
        print(f"[报告] 已保存: {report_path}")

# ========================= 4. 可视化 =========================
class Visualizer:
    """检测结果可视化工具"""

    @staticmethod
    def plot_metrics(results: dict, save_path: str = None):
        """绘制评估指标柱状图"""
        fig, axes = plt.subplots(1, 2, figsize=(12, 5))

        # 指标柱状图
        metrics = ['mAP50', 'mAP50-95', 'Precision', 'Recall', 'F1']
        values = [results.get(m, 0) for m in metrics]
        axes[0].bar(metrics, values, color=['#2E86AB', '#A23B72', '#F18F01', '#C73E1D', '#3B1F2B'])
        axes[0].set_ylim(0, 1)
        axes[0].set_title('Detection Metrics')
        axes[0].grid(axis='y', alpha=0.3)

        # 各类别AP（示意）
        # axes[1].barh(Config.CLASS_NAMES, ap_values)
        axes[1].set_title('Per-Class AP')
        axes[1].grid(axis='x', alpha=0.3)

        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300)
        plt.show()

# ========================= 5. 主程序 =========================
def main():
    print("=" * 60)
    print("  基于YOLO26的透明塑料瓶检测系统")
    print("=" * 60)

    config = Config()
    os.makedirs(config.PROJECT_DIR, exist_ok=True)
    os.makedirs(config.WEIGHT_DIR, exist_ok=True)
    os.makedirs(config.RESULT_DIR, exist_ok=True)

    # 检查数据集配置
    if not os.path.exists(config.DATA_YAML):
        print(f"[错误] 数据集配置文件不存在: {config.DATA_YAML}")
        print("请创建包含 train/val 路径及 nc/names 的 YAML 文件。")
        return

    # 初始化检测器
    detector = BottleDetector(config.MODEL_NAME)

    # 训练（如果已有权重则跳过）
    best_weight = os.path.join(config.WEIGHT_DIR, "best.pt")
    if not os.path.exists(best_weight):
        print("\n[提示] 未找到训练权重，开始训练...")
        detector.train(config.DATA_YAML, config)
        # 实际训练后可将最佳权重复制到 WEIGHT_DIR
    else:
        print(f"\n[提示] 加载已有权重: {best_weight}")
        detector.model = YOLO(best_weight)

    # 评估
    eval_results = detector.evaluate(config.DATA_YAML)

    # 可视化指标
    Visualizer.plot_metrics(eval_results, save_path=os.path.join(config.RESULT_DIR, "metrics.png"))

    # 批量推理示例
    test_image_dir = "./bottle_dataset/val/images"   # 替换为实际路径
    output_dir = os.path.join(config.RESULT_DIR, "predictions")
    if os.path.exists(test_image_dir):
        print("\n[推理] 对验证集进行批量检测...")
        detector.predict_batch(test_image_dir, output_dir, conf_thresh=0.25)

    print("\n" + "=" * 60)
    print("  系统运行完成！")
    print("=" * 60)

# ========================= 入口 =========================
if __name__ == "__main__":
    main()

# YOLOv8 / YOLOv11 Object Detection Configuration
path: D:/YOLO26/data/industrial-inspection-system
train: images/train
val: images/val

nc: 4
names:
  0: bottle
  1: cap
  2: label
  3: liquid

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

if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/26/yolo26.yaml')
    model.train(data='data/data_detect.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): 120 layers, 2,375,616 parameters, 0 gradients, 5.3 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% ━━━━━━━━━━━━ 1/1 4.1s/it 4.1s
                   all         24         83      0.976      0.964      0.965      0.904
                bottle         24         29      0.952          1      0.993      0.988
                   cap         18         18      0.959          1      0.995       0.98
                 label          7          8      0.993      0.875      0.875       0.76
                liquid         24         28          1      0.982      0.995      0.888

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