第7章 目标检测¶

import numpy as np

def calculate_iou(box1, box2):
    """计算两个边界框的IoU

    Args:
        box1: [x1, y1, x2, y2] 格式的边界框
        box2: [x1, y1, x2, y2] 格式的边界框
    Returns:
        iou: 交并比，范围 [0, 1]
    """
    # 计算交集区域
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])

    intersection = max(0, x2 - x1) * max(0, y2 - y1)

    # 计算各自面积
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])

    # 并集 = 两面积之和 - 交集
    union = area1 + area2 - intersection

    return intersection / union if union > 0 else 0.0

# 示例
pred_box = [50, 50, 200, 200]
gt_box   = [60, 40, 210, 190]
print(f"IoU = {calculate_iou(pred_box, gt_box):.4f}")
# IoU = 0.6806

def compute_ap(recalls, precisions):
    """计算单类别 AP（VOC 2007 的 11 点插值法）

    Args:
        recalls: 召回率列表（按置信度降序排列）
        precisions: 精确率列表（与 recalls 一一对应）
    Returns:
        ap: Average Precision
    """
    # 11 点插值
    ap = 0.0
    for t in np.arange(0, 1.1, 0.1):
        # 取 recall >= t 时的最大 precision
        precs = [p for r, p in zip(recalls, precisions) if r >= t]  # zip按位置配对
        if precs:
            ap += max(precs) / 11.0
    return ap

def compute_map(all_aps):
    """计算 mAP = 所有类别 AP 的均值"""
    return np.mean(all_aps)

# 使用 pycocotools 加载 COCO 数据集
from pycocotools.coco import COCO

coco = COCO('annotations/instances_val2017.json')
cat_ids = coco.getCatIds()         # 获取所有类别 id
img_ids = coco.getImgIds()         # 获取所有图像 id
print(f"COCO 类别数: {len(cat_ids)}, 图像数: {len(img_ids)}")

# 查看某张图像的标注
img_info = coco.loadImgs(img_ids[0])[0]
ann_ids = coco.getAnnIds(imgIds=img_info['id'])
anns = coco.loadAnns(ann_ids)
print(f"图像 {img_info['file_name']} 有 {len(anns)} 个标注")

import torch
import torch.nn as nn
import torchvision

class FastRCNN(nn.Module):  # 继承nn.Module定义网络层
    """Fast R-CNN 简化实现"""
    def __init__(self, num_classes=21):
        super().__init__()  # super()调用父类方法
        # Backbone: 共享卷积特征
        backbone = torchvision.models.vgg16(weights='DEFAULT')
        self.features = backbone.features       # 共享卷积层

        # RoI Pooling: 将任意大小区域转为 7×7
        self.roi_pool = torchvision.ops.RoIPool(
            output_size=(7, 7),
            spatial_scale=1.0 / 16  # VGG16 下采样 16 倍
        )

        # 分类 + 回归头
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
        )
        self.cls_head = nn.Linear(4096, num_classes)       # 类别
        self.reg_head = nn.Linear(4096, num_classes * 4)   # 边界框回归

    def forward(self, images, rois):
        """
        Args:
            images: (B, 3, H, W)
            rois: List[Tensor], 每个 Tensor shape = (N_i, 5)，
                  格式 [batch_idx, x1, y1, x2, y2]
        """
        feat = self.features(images)                       # 共享特征
        pooled = self.roi_pool(feat, rois)                 # (N, 512, 7, 7)
        pooled = pooled.flatten(1)                         # (N, 512*7*7)
        fc = self.classifier(pooled)                       # (N, 4096)
        cls_score = self.cls_head(fc)                      # (N, num_classes)
        bbox_pred = self.reg_head(fc)                      # (N, num_classes*4)
        return cls_score, bbox_pred

import torch
import torch.nn as nn
import torch.nn.functional as F

class RPN(nn.Module):
    """Region Proposal Network"""
    def __init__(self, in_channels=512, num_anchors=9):
        super().__init__()
        # 3×3 卷积提取局部特征
        self.conv = nn.Conv2d(in_channels, 512, 3, padding=1)
        # 分类头: 每个 anchor 的前景/背景二分类
        self.cls_head = nn.Conv2d(512, num_anchors * 2, 1)
        # 回归头: 每个 anchor 的边界框偏移 (dx, dy, dw, dh)
        self.reg_head = nn.Conv2d(512, num_anchors * 4, 1)

        # 权重初始化
        for layer in [self.conv, self.cls_head, self.reg_head]:
            nn.init.normal_(layer.weight, std=0.01)
            nn.init.constant_(layer.bias, 0)

    def forward(self, feature_map):
        """
        Args:
            feature_map: (B, C, H, W) 来自 backbone 的特征图
        Returns:
            cls_score: (B, num_anchors*2, H, W) 前景/背景分数
            bbox_pred: (B, num_anchors*4, H, W) 偏移量
        """
        x = F.relu(self.conv(feature_map))  # F.xxx PyTorch函数式API
        cls_score = self.cls_head(x)   # (B, 18, H, W)
        bbox_pred = self.reg_head(x)   # (B, 36, H, W)
        return cls_score, bbox_pred

def generate_anchors(feature_h, feature_w, stride=16):
    """生成所有 anchor 坐标

    Args:
        feature_h, feature_w: 特征图尺寸
        stride: 下采样步长
    Returns:
        anchors: (H*W*9, 4) 格式 [x1, y1, x2, y2]
    """
    scales = [128, 256, 512]          # anchor 尺度
    ratios = [0.5, 1.0, 2.0]         # 宽高比

    anchors = []
    for i in range(feature_h):
        for j in range(feature_w):
            cx = (j + 0.5) * stride   # anchor 中心 x
            cy = (i + 0.5) * stride   # anchor 中心 y
            for s in scales:
                for r in ratios:
                    w = s * (r ** 0.5)
                    h = s / (r ** 0.5)
                    anchors.append([
                        cx - w / 2, cy - h / 2,
                        cx + w / 2, cy + h / 2
                    ])
    return torch.tensor(anchors, dtype=torch.float32)

import torchvision.ops as ops

# RoI Pooling: 有量化误差（坐标取整）
roi_pool = ops.RoIPool(output_size=(7, 7), spatial_scale=1/16)

# RoI Align: 双线性插值，无量化误差，更精确
roi_align = ops.RoIAlign(
    output_size=(7, 7),
    spatial_scale=1.0 / 16,
    sampling_ratio=2          # 每个 bin 采样 2×2 个点
)

# 使用示例
feature_map = torch.randn(1, 256, 50, 50)   # backbone 输出
rois = torch.tensor([[0, 10.6, 20.3, 100.7, 150.2]])  # [batch_idx, x1, y1, x2, y2]

pooled = roi_align(feature_map, rois)  # (1, 256, 7, 7)

import torch
import torch.nn as nn
import torch.nn.functional as F

class FPN(nn.Module):
    """Feature Pyramid Network"""
    def __init__(self, in_channels_list, out_channels=256):
        """
        Args:
            in_channels_list: backbone 各层输出通道数，如 [256, 512, 1024, 2048]
            out_channels: FPN 统一输出通道数
        """
        super().__init__()
        # 横向连接: 1×1 卷积统一通道数
        self.lateral_convs = nn.ModuleList([
            nn.Conv2d(in_ch, out_channels, 1) for in_ch in in_channels_list
        ])
        # 输出卷积: 3×3 卷积消除上采样混叠
        self.output_convs = nn.ModuleList([
            nn.Conv2d(out_channels, out_channels, 3, padding=1)
            for _ in in_channels_list
        ])

    def forward(self, features):
        """
        Args:
            features: [C2, C3, C4, C5] backbone 多尺度特征
        Returns:
            fpn_features: [P2, P3, P4, P5]
        """
        # 先做横向连接
        laterals = [conv(f) for conv, f in zip(self.lateral_convs, features)]

        # 自顶向下融合
        for i in range(len(laterals) - 2, -1, -1):
            up = F.interpolate(laterals[i + 1], size=laterals[i].shape[2:], mode='nearest')
            laterals[i] = laterals[i] + up

        # 输出卷积
        outputs = [conv(lat) for conv, lat in zip(self.output_convs, laterals)]
        return outputs

# YOLOv8 使用示例（Ultralytics CLI & Python）
from ultralytics import YOLO

# --------- 推理 ---------
model = YOLO('yolov8n.pt')            # 加载预训练模型（nano 版）
results = model('bus.jpg')             # 推理
results[0].show()                      # 可视化

# --------- 训练 ---------
model = YOLO('yolov8s.pt')            # small 版本
model.train(
    data='coco128.yaml',              # 数据集配置
    epochs=100,
    imgsz=640,
    batch=16,
    device='0',                       # GPU
)

# --------- 导出 ---------
model.export(format='onnx')           # 导出 ONNX
model.export(format='engine')         # 导出 TensorRT

import torch
import torch.nn as nn

class FCOSHead(nn.Module):
    """FCOS 检测头（简化版）"""
    def __init__(self, in_channels=256, num_classes=80):
        super().__init__()
        # 分类分支
        cls_tower = []
        for _ in range(4):   # 4 层卷积
            cls_tower.append(nn.Conv2d(in_channels, in_channels, 3, padding=1))
            cls_tower.append(nn.GroupNorm(32, in_channels))
            cls_tower.append(nn.ReLU())
        self.cls_tower = nn.Sequential(*cls_tower)
        self.cls_logits = nn.Conv2d(in_channels, num_classes, 3, padding=1)

        # 回归分支
        reg_tower = []
        for _ in range(4):
            reg_tower.append(nn.Conv2d(in_channels, in_channels, 3, padding=1))
            reg_tower.append(nn.GroupNorm(32, in_channels))
            reg_tower.append(nn.ReLU())
        self.reg_tower = nn.Sequential(*reg_tower)
        self.bbox_pred = nn.Conv2d(in_channels, 4, 3, padding=1)  # (l, t, r, b)

        # center-ness 分支
        self.centerness = nn.Conv2d(in_channels, 1, 3, padding=1)

    def forward(self, x):
        cls_feat = self.cls_tower(x)
        reg_feat = self.reg_tower(x)

        cls_score = self.cls_logits(cls_feat)        # (B, C, H, W)
        bbox = torch.exp(self.bbox_pred(reg_feat))   # (B, 4, H, W) 正值
        ctr = self.centerness(cls_feat)               # (B, 1, H, W)
        return cls_score, bbox, ctr

class CenterNetHead(nn.Module):
    """CenterNet 检测头"""
    def __init__(self, in_channels=64, num_classes=80):
        super().__init__()
        # 热力图分支: 预测中心点
        self.heatmap = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels, num_classes, 1),
        )
        # 尺寸分支: 预测宽高
        self.wh = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels, 2, 1),   # (w, h)
        )
        # 偏移分支: 下采样导致的中心偏移
        self.offset = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels, 2, 1),   # (offset_x, offset_y)
        )

    def forward(self, x):
        hm = torch.sigmoid(self.heatmap(x))   # 热力图
        wh = self.wh(x)                        # 宽高
        off = self.offset(x)                   # 偏移
        return hm, wh, off

import torch
import torch.nn as nn
from scipy.optimize import linear_sum_assignment

class HungarianMatcher(nn.Module):
    """匈牙利匹配器: 预测与GT的最优一一匹配"""
    def __init__(self, cost_class=1.0, cost_bbox=5.0, cost_giou=2.0):
        super().__init__()
        self.cost_class = cost_class
        self.cost_bbox = cost_bbox
        self.cost_giou = cost_giou

    @torch.no_grad()  # 禁用梯度计算，节省内存
    def forward(self, outputs, targets):
        """
        Args:
            outputs: {"pred_logits": (B, N, C), "pred_boxes": (B, N, 4)}
            targets: [{"labels": (M,), "boxes": (M, 4)}, ...] per image
        Returns:
            indices: [(pred_idx, gt_idx), ...] 每张图的匹配结果
        """
        B, N = outputs["pred_logits"].shape[:2]  # 切片操作，取前n个元素

        # 拼接 batch
        out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1)  # (B*N, C)
        out_bbox = outputs["pred_boxes"].flatten(0, 1)               # (B*N, 4)

        tgt_ids = torch.cat([t["labels"] for t in targets])  # torch.cat沿已有维度拼接张量
        tgt_bbox = torch.cat([t["boxes"] for t in targets])

        # 计算代价矩阵
        cost_class = -out_prob[:, tgt_ids]                          # 分类代价
        cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)           # L1 代价
        cost_giou = -self._generalized_iou(out_bbox, tgt_bbox)     # GIoU 代价

        C = (self.cost_class * cost_class +
             self.cost_bbox * cost_bbox +
             self.cost_giou * cost_giou)
        C = C.view(B, N, -1).cpu()  # 重塑张量形状

        # 对每张图独立做匈牙利匹配
        sizes = [len(t["labels"]) for t in targets]
        indices = []
        for i, c in enumerate(C.split(sizes, -1)):  # enumerate同时获取索引和元素
            row_ind, col_ind = linear_sum_assignment(c[i])
            indices.append((
                torch.as_tensor(row_ind, dtype=torch.int64),
                torch.as_tensor(col_ind, dtype=torch.int64)
            ))
        return indices

    def _generalized_iou(self, boxes1, boxes2):
        """计算 Generalized IoU（简化版）"""
        # 简化实现：此处使用IoU替代GIoU。完整GIoU = IoU - (C_area - union) / C_area，其中C为最小外接矩形面积
        # 实际实现需要处理 cxcywh → xyxy 转换
        return torch.zeros(len(boxes1), len(boxes2))  # placeholder

class SimpleDETR(nn.Module):
    """DETR 简化版本（展示核心架构）"""
    def __init__(self, num_classes=91, num_queries=100, d_model=256):
        super().__init__()
        # Backbone: ResNet50（取最后一层特征）
        import torchvision
        backbone = torchvision.models.resnet50(weights='DEFAULT')
        self.backbone = nn.Sequential(*list(backbone.children())[:-2])

        # 降维
        self.input_proj = nn.Conv2d(2048, d_model, 1)

        # 位置编码
        # 注意：此处位置编码固定大小，实际DETR使用可学习或正弦位置编码，可适配不同输入尺寸
        self.pos_embed = nn.Parameter(torch.randn(1, d_model, 25, 25))

        # Transformer
        self.transformer = nn.Transformer(
            d_model=d_model, nhead=8,
            num_encoder_layers=6, num_decoder_layers=6,
            dim_feedforward=2048, dropout=0.1,
            batch_first=True
        )

        # Object Queries: 可学习的查询向量
        self.query_embed = nn.Embedding(num_queries, d_model)

        # 输出头
        self.class_head = nn.Linear(d_model, num_classes + 1)  # +1 = "no object"
        self.bbox_head = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.ReLU(),
            nn.Linear(d_model, d_model),
            nn.ReLU(),
            nn.Linear(d_model, 4),    # (cx, cy, w, h) 归一化坐标
            nn.Sigmoid()
        )

    def forward(self, x):
        # Backbone 特征
        feat = self.backbone(x)                     # (B, 2048, H/32, W/32)
        feat = self.input_proj(feat)                 # (B, 256, H', W')
        B, C, H, W = feat.shape

        # 展平 + 位置编码
        pos = self.pos_embed[:, :, :H, :W].flatten(2).permute(0, 2, 1)
        src = feat.flatten(2).permute(0, 2, 1) + pos  # (B, H'*W', 256)

        # Transformer
        queries = self.query_embed.weight.unsqueeze(0).repeat(B, 1, 1)  # (B, N, 256)  # unsqueeze增加一个维度
        hs = self.transformer(src, queries)           # (B, N, 256)

        # 输出
        cls = self.class_head(hs)                     # (B, N, num_classes+1)
        bbox = self.bbox_head(hs)                     # (B, N, 4)
        return {"pred_logits": cls, "pred_boxes": bbox}

# RT-DETR 使用 (Ultralytics 集成)
from ultralytics import RTDETR

model = RTDETR('rtdetr-l.pt')          # 加载预训练
results = model('bus.jpg')             # 推理
model.train(data='coco.yaml', epochs=100)  # 训练

import cv2
import numpy as np
import random

def mosaic_augmentation(images, labels_list, img_size=640):
    """Mosaic 数据增强: 4 张图拼接

    Args:
        images: 4 张图像列表
        labels_list: 4 个标注列表，每个元素 shape=(N, 5) [cls, cx, cy, w, h] 归一化
        img_size: 输出图像尺寸
    Returns:
        mosaic_img: 拼接后图像
        mosaic_labels: 合并后标注
    """
    s = img_size
    # 随机选择拼接中心点
    cx = int(random.uniform(s * 0.25, s * 0.75))
    cy = int(random.uniform(s * 0.25, s * 0.75))

    mosaic_img = np.zeros((s, s, 3), dtype=np.uint8)
    mosaic_labels = []

    # 4 个位置: 左上, 右上, 左下, 右下
    placements = [
        (0, 0, cx, cy),         # 左上
        (cx, 0, s, cy),         # 右上
        (0, cy, cx, s),         # 左下
        (cx, cy, s, s),         # 右下
    ]

    for i, (x1, y1, x2, y2) in enumerate(placements):
        img = images[i]
        h, w = img.shape[:2]

        # resize 到目标区域大小
        pw, ph = x2 - x1, y2 - y1
        resized = cv2.resize(img, (pw, ph))
        mosaic_img[y1:y2, x1:x2] = resized

        # 转换标注坐标
        if len(labels_list[i]) > 0:
            labels = labels_list[i].copy()
            labels[:, 1] = labels[:, 1] * pw / s + x1 / s  # cx
            labels[:, 2] = labels[:, 2] * ph / s + y1 / s  # cy
            labels[:, 3] = labels[:, 3] * pw / s            # w
            labels[:, 4] = labels[:, 4] * ph / s            # h
            mosaic_labels.append(labels)

    if mosaic_labels:
        mosaic_labels = np.concatenate(mosaic_labels, axis=0)
    else:
        mosaic_labels = np.zeros((0, 5))

    return mosaic_img, mosaic_labels

def mixup_augmentation(img1, labels1, img2, labels2, alpha=1.5):
    """MixUp 数据增强"""
    lam = np.random.beta(alpha, alpha)

    # 确保尺寸相同
    h = max(img1.shape[0], img2.shape[0])
    w = max(img1.shape[1], img2.shape[1])
    img1 = cv2.resize(img1, (w, h))
    img2 = cv2.resize(img2, (w, h))

    mixed_img = (lam * img1 + (1 - lam) * img2).astype(np.uint8)
    mixed_labels = np.concatenate([labels1, labels2], axis=0)

    return mixed_img, mixed_labels

import torch
import torch.nn.functional as F

def focal_loss(pred, target, alpha=0.25, gamma=2.0):
    """Focal Loss 实现

    Args:
        pred: (N, C) 预测 logits
        target: (N,) 目标类别
        alpha: 平衡因子
        gamma: 聚焦参数
    """
    ce_loss = F.cross_entropy(pred, target, reduction='none')  # F.cross_entropy PyTorch函数式交叉熵损失
    p_t = torch.exp(-ce_loss)              # 预测正确的概率
    focal_weight = alpha * (1 - p_t) ** gamma
    loss = focal_weight * ce_loss
    return loss.mean()

import torch

def ciou_loss(pred_boxes, target_boxes):
    """CIoU Loss 实现

    Args:
        pred_boxes: (N, 4) 格式 [x1, y1, x2, y2]
        target_boxes: (N, 4) 格式 [x1, y1, x2, y2]
    Returns:
        loss: CIoU Loss 均值
    """
    # 交集
    inter_x1 = torch.max(pred_boxes[:, 0], target_boxes[:, 0])
    inter_y1 = torch.max(pred_boxes[:, 1], target_boxes[:, 1])
    inter_x2 = torch.min(pred_boxes[:, 2], target_boxes[:, 2])
    inter_y2 = torch.min(pred_boxes[:, 3], target_boxes[:, 3])
    inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * \
                 torch.clamp(inter_y2 - inter_y1, min=0)

    # 各自面积
    area1 = (pred_boxes[:, 2] - pred_boxes[:, 0]) * (pred_boxes[:, 3] - pred_boxes[:, 1])
    area2 = (target_boxes[:, 2] - target_boxes[:, 0]) * (target_boxes[:, 3] - target_boxes[:, 1])
    union = area1 + area2 - inter_area
    iou = inter_area / (union + 1e-7)

    # 最小外接矩形对角线
    enclose_x1 = torch.min(pred_boxes[:, 0], target_boxes[:, 0])
    enclose_y1 = torch.min(pred_boxes[:, 1], target_boxes[:, 1])
    enclose_x2 = torch.max(pred_boxes[:, 2], target_boxes[:, 2])
    enclose_y2 = torch.max(pred_boxes[:, 3], target_boxes[:, 3])
    c2 = (enclose_x2 - enclose_x1) ** 2 + (enclose_y2 - enclose_y1) ** 2 + 1e-7

    # 中心点距离
    pred_cx = (pred_boxes[:, 0] + pred_boxes[:, 2]) / 2
    pred_cy = (pred_boxes[:, 1] + pred_boxes[:, 3]) / 2
    gt_cx = (target_boxes[:, 0] + target_boxes[:, 2]) / 2
    gt_cy = (target_boxes[:, 1] + target_boxes[:, 3]) / 2
    d2 = (pred_cx - gt_cx) ** 2 + (pred_cy - gt_cy) ** 2

    # 宽高比一致性
    pred_w = pred_boxes[:, 2] - pred_boxes[:, 0]
    pred_h = pred_boxes[:, 3] - pred_boxes[:, 1]
    gt_w = target_boxes[:, 2] - target_boxes[:, 0]
    gt_h = target_boxes[:, 3] - target_boxes[:, 1]

    v = (4 / (torch.pi ** 2)) * (torch.atan(gt_w / (gt_h + 1e-7)) -
                                   torch.atan(pred_w / (pred_h + 1e-7))) ** 2
    with torch.no_grad():
        alpha = v / (1 - iou + v + 1e-7)

    ciou = iou - d2 / c2 - alpha * v
    loss = 1 - ciou
    return loss.mean()

def nms(boxes, scores, iou_threshold=0.5):
    """非极大值抑制

    Args:
        boxes: (N, 4) 边界框 [x1, y1, x2, y2]
        scores: (N,) 置信度分数
        iou_threshold: IoU 阈值
    Returns:
        keep: 保留的框索引
    """
    order = scores.argsort()[::-1]    # 按分数降序排列
    keep = []

    while len(order) > 0:
        i = order[0]
        keep.append(i)

        if len(order) == 1:
            break

        # 计算当前最高分框与其余框的 IoU
        ious = np.array([  # np.array创建NumPy数组
            calculate_iou(boxes[i], boxes[j]) for j in order[1:]
        ])

        # 保留 IoU 小于阈值的框
        mask = ious < iou_threshold
        order = order[1:][mask]

    return keep

def soft_nms(boxes, scores, sigma=0.5, score_threshold=0.001):
    """Soft-NMS 实现"""
    N = len(boxes)
    indices = np.arange(N)

    for i in range(N):
        # 找当前最高分
        max_idx = np.argmax(scores[i:]) + i
        # 交换到当前位置
        boxes[[i, max_idx]] = boxes[[max_idx, i]]
        scores[[i, max_idx]] = scores[[max_idx, i]]
        indices[[i, max_idx]] = indices[[max_idx, i]]

        # 衰减重叠框的分数
        for j in range(i + 1, N):
            iou = calculate_iou(boxes[i], boxes[j])
            scores[j] *= np.exp(-(iou ** 2) / sigma)  # Gaussian 衰减

    keep = np.where(scores > score_threshold)[0]
    return indices[keep]

# SAHI 切片推理（pip install sahi）
from sahi import AutoDetectionModel
from sahi.predict import get_sliced_prediction

# 加载检测模型
detection_model = AutoDetectionModel.from_pretrained(
    model_type='yolov8',
    model_path='yolov8s.pt',
    confidence_threshold=0.3,
    device='cuda:0',
)

# 切片推理
result = get_sliced_prediction(
    image='high_res_image.jpg',
    detection_model=detection_model,
    slice_height=640,           # 切片高度
    slice_width=640,            # 切片宽度
    overlap_height_ratio=0.2,   # 高度方向重叠比例
    overlap_width_ratio=0.2,    # 宽度方向重叠比例
    perform_standard_pred=True, # 同时做全图推理
    postprocess_type='NMM',     # 非极大合并
    postprocess_match_threshold=0.5,
)

# 可视化
result.export_visuals(export_dir='results/')
print(f"检测到 {len(result.object_prediction_list)} 个目标")

"""
YOLOv8 安全帽检测 — 完整训练脚本
"""
from ultralytics import YOLO

def train():
    """训练模型"""
    # 1. 加载预训练模型（迁移学习）
    model = YOLO('yolov8s.pt')    # small 版，平衡精度与速度

    # 2. 开始训练
    results = model.train(
        data='helmet.yaml',       # 数据集配置
        epochs=100,               # 训练轮数
        imgsz=640,                # 输入尺寸
        batch=16,                 # batch size
        device='0',               # GPU 设备号
        workers=8,                # 数据加载线程数

        # 优化器与学习率
        optimizer='SGD',          # 优化器
        lr0=0.01,                 # 初始学习率
        lrf=0.01,                 # 最终学习率 = lr0 * lrf
        momentum=0.937,           # SGD 动量
        weight_decay=0.0005,      # 权重衰减
        warmup_epochs=3,          # warmup 轮数

        # 数据增强
        mosaic=1.0,               # Mosaic 概率（最后 10 个 epoch 关闭）
        mixup=0.1,                # MixUp 概率
        copy_paste=0.1,           # CopyPaste 概率
        hsv_h=0.015,              # 色调增强
        hsv_s=0.7,                # 饱和度增强
        hsv_v=0.4,                # 明度增强
        degrees=0.0,              # 旋转角度
        translate=0.1,            # 平移比例
        scale=0.5,                # 缩放比例
        fliplr=0.5,               # 水平翻转概率

        # 保存与日志
        project='runs/helmet',    # 项目目录
        name='exp',               # 实验名称
        save=True,                # 保存 checkpoint
        save_period=10,           # 每 10 个 epoch 保存
        val=True,                 # 训练时验证
        plots=True,               # 生成训练曲线
    )
    return results

def validate():
    """验证模型"""
    model = YOLO('runs/helmet/exp/weights/best.pt')
    metrics = model.val(
        data='helmet.yaml',
        imgsz=640,
        batch=16,
        conf=0.25,
        iou=0.6,
        device='0',
    )

    print(f"mAP50    : {metrics.box.map50:.4f}")
    print(f"mAP50-95 : {metrics.box.map:.4f}")
    # 各类别 AP
    for i, name in enumerate(metrics.names.values()):
        print(f"  {name:15s}: AP50={metrics.box.ap50[i]:.4f}")

    return metrics

if __name__ == '__main__':
    train()
    validate()

from ultralytics import YOLO
from PIL import Image
import cv2

def inference_single(image_path):
    """单张图片推理"""
    model = YOLO('runs/helmet/exp/weights/best.pt')

    results = model(image_path, conf=0.3, iou=0.5)

    # 解析结果
    for r in results:
        boxes = r.boxes
        for box in boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            xyxy = box.xyxy[0].cpu().numpy().astype(int)
            label = r.names[cls_id]
            print(f"[{label}] conf={conf:.2f} box={xyxy}")

    # 保存可视化
    annotated = results[0].plot()          # numpy (BGR)
    cv2.imwrite('result.jpg', annotated)
    print("结果已保存至 result.jpg")

def inference_video(video_path, output_path='output.mp4'):
    """视频推理"""
    model = YOLO('runs/helmet/exp/weights/best.pt')

    cap = cv2.VideoCapture(video_path)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

    frame_count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        results = model(frame, conf=0.3, verbose=False)
        annotated = results[0].plot()
        writer.write(annotated)

        frame_count += 1
        if frame_count % 100 == 0:
            print(f"已处理 {frame_count} 帧")

    cap.release()
    writer.release()
    print(f"视频结果已保存至 {output_path}")

from ultralytics import YOLO

model = YOLO('runs/helmet/exp/weights/best.pt')

# 导出 ONNX（通用部署格式）
model.export(
    format='onnx',
    imgsz=640,
    simplify=True,         # 简化 ONNX 图
    opset=12,              # ONNX opset 版本
    dynamic=False,         # 固定输入尺寸（部署更快）
)

# 导出 TensorRT（NVIDIA GPU 高性能推理）
model.export(
    format='engine',
    imgsz=640,
    half=True,             # FP16 量化
    device='0',
)

# 导出 NCNN（移动端）
model.export(format='ncnn', imgsz=640)

# 导出 CoreML（iOS）
model.export(format='coreml', imgsz=640)

import onnxruntime as ort
import numpy as np
import cv2

class YOLOv8ONNXDetector:
    """YOLOv8 ONNX 部署推理器"""

    def __init__(self, model_path, conf_thres=0.3, iou_thres=0.5):
        # 创建推理会话
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        self.session = ort.InferenceSession(model_path, providers=providers)

        self.input_name = self.session.get_inputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape   # [1, 3, 640, 640]
        self.imgsz = self.input_shape[2]

        self.conf_thres = conf_thres
        self.iou_thres = iou_thres

    def preprocess(self, image):
        """预处理：letterbox resize + 归一化"""
        h, w = image.shape[:2]
        scale = min(self.imgsz / h, self.imgsz / w)
        new_h, new_w = int(h * scale), int(w * scale)

        # Resize
        resized = cv2.resize(image, (new_w, new_h))

        # Letterbox padding
        canvas = np.full((self.imgsz, self.imgsz, 3), 114, dtype=np.uint8)
        dh, dw = (self.imgsz - new_h) // 2, (self.imgsz - new_w) // 2
        canvas[dh:dh + new_h, dw:dw + new_w] = resized

        # BGR → RGB, HWC → CHW, 归一化
        blob = canvas[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) / 255.0
        blob = np.expand_dims(blob, axis=0)    # (1, 3, 640, 640)

        return blob, scale, dw, dh

    def postprocess(self, output, scale, dw, dh):
        """后处理：解码 + NMS"""
        # YOLOv8 输出: (1, 4+num_classes, 8400) → 转置为 (8400, 4+C)
        preds = output[0].squeeze(0).T   # (8400, 4+C)  # squeeze压缩维度

        # 提取类别分数
        boxes = preds[:, :4]             # (8400, 4) cxcywh
        scores = preds[:, 4:]            # (8400, C)

        # 获取每个框的最大类别分数
        max_scores = scores.max(axis=1)
        class_ids = scores.argmax(axis=1)

        # 置信度过滤
        mask = max_scores > self.conf_thres
        boxes = boxes[mask]
        max_scores = max_scores[mask]
        class_ids = class_ids[mask]

        if len(boxes) == 0:
            return [], [], []

        # cxcywh → xyxy
        x1 = boxes[:, 0] - boxes[:, 2] / 2
        y1 = boxes[:, 1] - boxes[:, 3] / 2
        x2 = boxes[:, 0] + boxes[:, 2] / 2
        y2 = boxes[:, 1] + boxes[:, 3] / 2
        xyxy = np.stack([x1, y1, x2, y2], axis=1)

        # 还原到原图坐标
        xyxy[:, [0, 2]] = (xyxy[:, [0, 2]] - dw) / scale
        xyxy[:, [1, 3]] = (xyxy[:, [1, 3]] - dh) / scale

        # NMS
        keep = cv2.dnn.NMSBoxes(
            xyxy.tolist(),
            max_scores.tolist(),
            self.conf_thres,
            self.iou_thres,
        )

        if len(keep) > 0:
            keep = keep.flatten()
            return xyxy[keep], max_scores[keep], class_ids[keep]
        return [], [], []

    def detect(self, image):
        """完整检测流程"""
        blob, scale, dw, dh = self.preprocess(image)
        output = self.session.run(None, {self.input_name: blob})
        boxes, scores, class_ids = self.postprocess(output, scale, dw, dh)
        return boxes, scores, class_ids

    def draw(self, image, boxes, scores, class_ids, class_names):
        """绘制检测结果"""
        colors = [(0, 255, 0), (0, 0, 255), (255, 0, 0)]
        for box, score, cid in zip(boxes, scores, class_ids):
            x1, y1, x2, y2 = box.astype(int)
            color = colors[int(cid) % len(colors)]
            label = f"{class_names[int(cid)]} {score:.2f}"
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            cv2.putText(image, label, (x1, y1 - 8),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
        return image

# ============ 使用示例 ============
if __name__ == '__main__':
    detector = YOLOv8ONNXDetector('best.onnx', conf_thres=0.3, iou_thres=0.5)

    image = cv2.imread('test.jpg')
    boxes, scores, class_ids = detector.detect(image)

    class_names = ['helmet', 'no_helmet', 'person']
    result = detector.draw(image, boxes, scores, class_ids, class_names)
    cv2.imwrite('onnx_result.jpg', result)
    print(f"ONNX 推理完成，检测到 {len(boxes)} 个目标")