基于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

# -*- coding: utf-8 -*-
"""
基于 YOLO26 的箱子和仓库检测系统 —— 伪代码 / 可运行模板

算法伪代码：
输入：图像/视频/摄像头流 I，YOLO26 权重 W，置信度阈值 tau_conf，IoU 阈值 tau_iou
输出：检测框、类别、置信度、箱子数量 N_box、仓库数量 N_warehouse

1. model = YOLO(W)
2. for frame in I:
3.     result = model.predict(frame, conf=tau_conf, iou=tau_iou)
4.     for det in result.boxes:
5.         if det.cls == 0: N_box += 1
6.         if det.cls == 1: N_warehouse += 1
7.         draw_box(frame, det)
8.     save_or_show(frame)
9. return detections, N_box, N_warehouse

数据集类别：
0: box       箱子
1: warehouse 仓库
"""

import argparse
from pathlib import Path

import cv2
from ultralytics import YOLO

# =========================
# 全局配置
# =========================
CLASS_NAMES = {
    0: "box",
    1: "warehouse",
}

CLASS_NAMES_CN = {
    0: "箱子",
    1: "仓库",
}

COLORS = {
    0: (0, 255, 0),      # 箱子：绿色
    1: (255, 128, 0),    # 仓库：橙色
}

DATASET_YAML_TEMPLATE = """# datasets/box_warehouse.yaml
path: ./datasets/box_warehouse
train: images/train
val: images/val
test: images/test

names:
  0: box
  1: warehouse
"""


# =========================
# 数据集 YAML 模板
# =========================
def create_dataset_yaml(out_path: str = "datasets/box_warehouse.yaml"):
    """生成数据集配置文件模板。"""
    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(DATASET_YAML_TEMPLATE, encoding="utf-8")
    print(f"已生成数据集配置模板: {out}")


# =========================
# 训练
# =========================
def train(args):
    """训练 YOLO26 箱子和仓库检测模型。"""
    model = YOLO(args.model)

    model.train(
        data=args.data,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device,
        workers=args.workers,
        project=args.project,
        name=args.name,
        exist_ok=True,
        optimizer="auto",
        patience=50,
        seed=0,
        plots=True,
    )

    print("训练完成。")


# =========================
# 验证
# =========================
def val(args):
    """验证模型性能。"""
    model = YOLO(args.weights)

    metrics = model.val(
        data=args.data,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device,
        conf=0.001,
        iou=0.6,
        plots=True,
    )

    print(f"mAP50:    {metrics.box.map50:.4f}")
    print(f"mAP50-95: {metrics.box.map:.4f}")


# =========================
# 检测结果统计
# =========================
def count_objects(result):
    """统计当前帧中箱子和仓库数量。"""
    counts = {name: 0 for name in CLASS_NAMES.values()}

    boxes = result.boxes
    if boxes is None:
        return counts

    for cls_id in boxes.cls.cpu().numpy().astype(int):
        name = CLASS_NAMES.get(cls_id, str(cls_id))
        counts[name] = counts.get(name, 0) + 1

    return counts


# =========================
# 检测框绘制
# =========================
def draw_boxes(result):
    """在图像上绘制检测框和标签。"""
    img = result.orig_img.copy()
    boxes = result.boxes

    if boxes is None:
        return img

    xyxy = boxes.xyxy.cpu().numpy()
    confs = boxes.conf.cpu().numpy()
    clss = boxes.cls.cpu().numpy().astype(int)

    for (x1, y1, x2, y2), conf, cls_id in zip(xyxy, confs, clss):
        x1, y1, x2, y2 = map(int, (x1, y1, x2, y2))

        name = CLASS_NAMES.get(cls_id, str(cls_id))
        color = COLORS.get(cls_id, (0, 255, 0))

        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)

        label = f"{name} {conf:.2f}"
        (tw, th), _ = cv2.getTextSize(
            label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2
        )

        cv2.rectangle(
            img,
            (x1, y1 - th - 8),
            (x1 + tw, y1),
            color,
            -1,
        )
        cv2.putText(
            img,
            label,
            (x1, y1 - 5),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (255, 255, 255),
            2,
        )

    return img


# =========================
# 推理 / 检测
# =========================
def predict(args):
    """对图片、视频、目录或摄像头进行检测。"""
    model = YOLO(args.weights)

    # 如果 source 是 "0"，则使用摄像头
    source = int(args.source) if str(args.source).isdigit() else args.source

    results = model.predict(
        source=source,
        imgsz=args.imgsz,
        conf=args.conf,
        iou=args.iou,
        device=args.device,
        stream=True,
        save=args.save,
        project=args.project,
        name=args.name,
        verbose=False,
    )

    for i, result in enumerate(results):
        counts = count_objects(result)

        print(
            f"Frame {i}: "
            f"box={counts.get('box', 0)}, "
            f"warehouse={counts.get('warehouse', 0)}"
        )

        if args.show:
            annotated = draw_boxes(result)
            cv2.imshow("YOLO26 Box & Warehouse Detection", annotated)

            if cv2.waitKey(1) & 0xFF == ord("q"):
                break

    cv2.destroyAllWindows()


# =========================
# 模型导出
# =========================
def export(args):
    """导出模型为 ONNX、TorchScript、TensorRT 等格式。"""
    model = YOLO(args.weights)
    model.export(
        format=args.format,
        imgsz=args.imgsz,
        device=args.device,
    )
    print(f"模型已导出为 {args.format} 格式。")


# =========================
# 命令行参数
# =========================
def build_parser():
    parser = argparse.ArgumentParser(
        description="基于 YOLO26 的箱子和仓库检测系统"
    )
    sub = parser.add_subparsers(dest="mode", required=True)

    # 训练
    p_train = sub.add_parser("train", help="训练 YOLO26 模型")
    p_train.add_argument(
        "--model",
        default="yolo26n.pt",
        help="YOLO26 权重或配置，例如 yolo26n.pt",
    )
    p_train.add_argument(
        "--data",
        default="datasets/box_warehouse.yaml",
        help="数据集 YAML 文件",
    )
    p_train.add_argument("--epochs", type=int, default=100)
    p_train.add_argument("--imgsz", type=int, default=640)
    p_train.add_argument("--batch", type=int, default=16)
    p_train.add_argument("--device", default="0", help="cuda 设备号或 cpu")
    p_train.add_argument("--workers", type=int, default=8)
    p_train.add_argument("--project", default="runs/train")
    p_train.add_argument("--name", default="yolo26_box_warehouse")
    p_train.set_defaults(func=train)

    # 验证
    p_val = sub.add_parser("val", help="验证模型")
    p_val.add_argument(
        "--weights",
        default="runs/train/yolo26_box_warehouse/weights/best.pt",
    )
    p_val.add_argument(
        "--data",
        default="datasets/box_warehouse.yaml",
    )
    p_val.add_argument("--imgsz", type=int, default=640)
    p_val.add_argument("--batch", type=int, default=16)
    p_val.add_argument("--device", default="0")
    p_val.set_defaults(func=val)

    # 推理
    p_pred = sub.add_parser("predict", help="推理检测")
    p_pred.add_argument(
        "--weights",
        default="runs/train/yolo26_box_warehouse/weights/best.pt",
    )
    p_pred.add_argument(
        "--source",
        default="datasets/box_warehouse/images/val",
        help="图片、视频、目录或摄像头编号 0",
    )
    p_pred.add_argument("--imgsz", type=int, default=640)
    p_pred.add_argument("--conf", type=float, default=0.25)
    p_pred.add_argument("--iou", type=float, default=0.45)
    p_pred.add_argument("--device", default="0")
    p_pred.add_argument("--save", action="store_true", help="保存检测结果")
    p_pred.add_argument("--show", action="store_true", help="显示检测窗口")
    p_pred.add_argument("--project", default="runs/detect")
    p_pred.add_argument("--name", default="yolo26_box_warehouse_predict")
    p_pred.set_defaults(func=predict)

    # 导出
    p_export = sub.add_parser("export", help="导出模型")
    p_export.add_argument(
        "--weights",
        default="runs/train/yolo26_box_warehouse/weights/best.pt",
    )
    p_export.add_argument(
        "--format",
        default="onnx",
        help="onnx / torchscript / engine 等",
    )
    p_export.add_argument("--imgsz", type=int, default=640)
    p_export.add_argument("--device", default="0")
    p_export.set_defaults(func=export)

    # 生成数据集 YAML 模板
    p_data = sub.add_parser("make-yaml", help="生成数据集 YAML 模板")
    p_data.add_argument(
        "--out",
        default="datasets/box_warehouse.yaml",
    )
    p_data.set_defaults(func=lambda a: create_dataset_yaml(a.out))

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()

# Dataset unificado a partir de Final_dataset.
path: D:/YOLO26/data/UnifiedDataset

train: train/images
val: val/images
test: test/images

nc: 1
names: ["bulto"]

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

if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/26/yolo26.yaml')
    model.train(data='data/UnifiedDataset.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,375,031 parameters, 0 gradients, 5.3 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% ━━━━━━━━━━━━ 3/3 2.2it/s 1.4s
                   all         88        246       0.93      0.925      0.969      0.876

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