目标检测：从 R-CNN 到 YOLO

import torch

def compute_iou(box1: torch.Tensor, box2: torch.Tensor) -> float:
    """计算两个边界框的 IoU
    box格式: (x_center, y_center, w, h)
    """
    # 转换为 (x_min, y_min, x_max, y_max)
    b1_x1 = box1[0] - box1[2] / 2
    b1_y1 = box1[1] - box1[3] / 2
    b1_x2 = box1[0] + box1[2] / 2
    b1_y2 = box1[1] + box1[3] / 2

    b2_x1 = box2[0] - box2[2] / 2
    b2_y1 = box2[1] - box2[3] / 2
    b2_x2 = box2[0] + box2[2] / 2
    b2_y2 = box2[1] + box2[3] / 2

    # 交集
    inter_x1 = max(b1_x1, b2_x1)
    inter_y1 = max(b1_y1, b2_y1)
    inter_x2 = min(b1_x2, b2_x2)
    inter_y2 = min(b1_y2, b2_y2)
    inter_area = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)

    # 并集
    area1 = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    area2 = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
    union_area = area1 + area2 - inter_area

    return inter_area / union_area if union_area > 0 else 0.0

import torch
import torch.nn as nn
import torchvision

class FasterRCNNDemo(nn.Module):
    """Faster R-CNN 简化结构演示"""
    def __init__(self, num_classes: int = 80):
        super().__init__()
        # Backbone: ResNet-50 + FPN
        backbone = torchvision.models.resnet50(weights="IMAGENET1K_V1")
        backbone = nn.Sequential(*list(backbone.children())[:-2])

        self.backbone = backbone
        self.rpn = nn.Conv2d(2048, 512, kernel_size=3, padding=1)
        self.cls_head = nn.Sequential(
            nn.AdaptiveAvgPool2d((7, 7)),
            nn.Flatten(),
            nn.Linear(2048, 1024),
            nn.ReLU(),
            nn.Linear(1024, num_classes + 1),  # +1 for background
        )
        self.bbox_head = nn.Sequential(
            nn.AdaptiveAvgPool2d((7, 7)),
            nn.Flatten(),
            nn.Linear(2048, 1024),
            nn.ReLU(),
            nn.Linear(1024, num_classes * 4),  # 4 offsets per class
        )

    def forward(self, x):
        features = self.backbone(x)  # [B, 2048, H/32, W/32]
        rpn_scores = self.rpn(features)
        # ... RPN 生成 proposals → RoI Align → 分类/回归
        return rpn_scores

import torch
import torch.nn as nn

class YOLOv1Head(nn.Module):
    """YOLOv1 检测头简化实现"""
    def __init__(self, num_classes: int = 20, num_boxes: int = 2, grid_size: int = 7):
        super().__init__()
        self.num_classes = num_classes
        self.num_boxes = num_boxes
        self.grid_size = grid_size
        # 输出: S x S x (B*5 + C)
        output_dim = grid_size * grid_size * (num_boxes * 5 + num_classes)
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d((grid_size, grid_size)),
            nn.Flatten(),
            nn.Linear(1024 * grid_size * grid_size, 4096),
            nn.LeakyReLU(0.1),
            nn.Linear(4096, output_dim),
        )

    def forward(self, x):
        """x: [B, 1024, H, W] 来自 backbone"""
        output = self.head(x)  # [B, S*S*(B*5+C)]
        output = output.view(
            -1, self.grid_size, self.grid_size,
            self.num_boxes * 5 + self.num_classes
        )
        return output

    def decode_predictions(self, output, conf_threshold=0.5):
        """解析输出张量为检测框"""
        batch_size, S, _, _ = output.shape
        box_dim = 5
        confidences = output[:, :, :, 4::box_dim]  # 提取置信度
        return confidences

# YOLOv8 Anchor-Free 检测头 vs YOLOv5 Anchor-Based
import torch
import torch.nn as nn

class YOLOv8_Detect(nn.Module):
    """YOLOv8 Anchor-Free 检测头"""
    def __init__(self, num_classes: int = 80, reg_max: int = 16):
        super().__init__()
        self.num_classes = num_classes
        self.reg_max = reg_max
        # 解耦头：分类和回归分开
        self.cls_convs = nn.Conv2d(256, 256, 3, padding=1)
        self.reg_convs = nn.Conv2d(256, 256, 3, padding=1)
        self.cls_pred = nn.Conv2d(256, num_classes, 1)
        self.reg_pred = nn.Conv2d(256, 4 * (reg_max + 1), 1)  # DFLLoss

    def forward(self, x):
        """x: 多尺度特征图列表 [P3, P4, P5]"""
        results = []
        for feat in x:
            cls_feat = self.cls_convs(feat)
            reg_feat = self.reg_convs(feat)
            cls_out = self.cls_pred(cls_feat)
            reg_out = self.reg_pred(reg_feat)
            results.append(torch.cat([reg_out, cls_out], dim=1))
        return results

# YOLOv5/v8 推理示例（使用 Ultralytics）
# from ultralytics import YOLO
# model = YOLO("yolov8n.pt")  # 加载预训练 Nano 模型
# results = model("image.jpg")  # 推理
# results[0].show()  # 可视化

import torch
import torch.nn as nn
import math

class SSD300(nn.Module):
    """SSD300 简化实现 - 多尺度检测"""
    def __init__(self, num_classes: int = 21):
        super().__init__()
        self.num_classes = num_classes

        # Backbone: VGG-16 前30层
        self.base = nn.Sequential(
            *list(torchvision.models.vgg16(weights="IMAGENET1K_V1")
                  .features.children())[:30]
        )

        # 额外卷积层 (逐渐缩小特征图)
        self.extra = nn.ModuleList([
            nn.Conv2d(512, 256, 1),
            nn.Conv2d(256, 512, 3, stride=2, padding=1),
            nn.Conv2d(512, 128, 1),
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
        ])

        # 不同尺度特征图上的检测头
        # 38x38: 4 boxes/pixel, 19x19: 6 boxes/pixel, etc.
        self.loc = nn.ModuleList([
            nn.Conv2d(512, 4 * 4, 3, padding=1),   # 38x38
            nn.Conv2d(1024, 6 * 4, 3, padding=1),  # 19x19
            nn.Conv2d(512, 6 * 4, 3, padding=1),   # 10x10
            nn.Conv2d(256, 6 * 4, 3, padding=1),   # 5x5
            nn.Conv2d(256, 4 * 4, 3, padding=1),   # 3x3
            nn.Conv2d(256, 4 * 4, 3, padding=1),   # 1x1
        ])
        self.conf = nn.ModuleList([
            nn.Conv2d(512, num_classes * 4, 3, padding=1),
            nn.Conv2d(1024, num_classes * 6, 3, padding=1),
            nn.Conv2d(512, num_classes * 6, 3, padding=1),
            nn.Conv2d(256, num_classes * 6, 3, padding=1),
            nn.Conv2d(256, num_classes * 4, 3, padding=1),
            nn.Conv2d(256, num_classes * 4, 3, padding=1),
        ])

    def forward(self, x):
        # 多尺度特征提取 + 检测
        loc_preds, conf_preds = [], []
        # ... 前向传播
        return loc_preds, conf_preds

import torch

def nms(boxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float = 0.5):
    """标准 NMS 实现
    boxes: [N, 4] (x_min, y_min, x_max, y_max)
    scores: [N]
    返回: 保留框的索引
    """
    areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    _, order = scores.sort(descending=True)
    keep = []

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

        if order.numel() == 1:
            break

        # 计算 IoU
        xx1 = boxes[order[1:], 0].clamp(min=boxes[i, 0])
        yy1 = boxes[order[1:], 1].clamp(min=boxes[i, 1])
        xx2 = boxes[order[1:], 2].clamp(max=boxes[i, 2])
        yy2 = boxes[order[1:], 3].clamp(max=boxes[i, 3])

        inter = (xx2 - xx1).clamp(min=0) * (yy2 - yy1).clamp(min=0)
        union = areas[i] + areas[order[1:]] - inter
        iou = inter / union

        # 保留 IoU < 阈值的框
        ids = (iou < iou_threshold).nonzero(as_tuple=False).squeeze()
        if ids.numel() == 0:
            break
        order = order[ids + 1]

    return torch.tensor(keep, dtype=torch.long)

def soft_nms(boxes, scores, sigma=0.5, Nt=0.3, threshold=0.001):
    """Soft-NMS 高斯衰减版本"""
    N = scores.shape[0]
    for i in range(N):
        max_idx = i + scores[i:].argmax()
        scores[[i, max_idx]] = scores[[max_idx, i]]
        boxes[[i, max_idx]] = boxes[[max_idx, i]]

        for j in range(i + 1, N):
            iou = compute_iou(boxes[i], boxes[j])
            weight = torch.exp(-(iou * iou) / sigma)
            scores[j] *= weight

    return boxes[scores > threshold], scores[scores > threshold]

import cv2
import torch
import numpy as np
from pathlib import Path

class YOLODetector:
    """基于 PyTorch + YOLO 的目标检测器"""

    def __init__(self, model_name="yolov5s", conf_thres=0.25, iou_thres=0.45):
        self.model = torch.hub.load(
            "ultralytics/yolov5", model_name, pretrained=True
        )
        self.model.eval()
        self.classes = self.model.names
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres

    def detect(self, image_path):
        """对单张图像进行目标检测"""
        # 推理
        results = self.model(image_path)
        # 解析结果 (xyxy, conf, cls)
        detections = results.pred[0]
        boxes = detections[:, :4]
        confidences = detections[:, 4]
        class_ids = detections[:, 5].int()

        return boxes, confidences, class_ids

    def visualize(self, image, boxes, confidences, class_ids):
        """可视化检测结果"""
        colors = np.random.randint(0, 255, size=(len(self.classes), 3))

        for box, conf, cls_id in zip(boxes, confidences, class_ids):
            x1, y1, x2, y2 = map(int, box)
            color = colors[cls_id].tolist()
            label = f"{self.classes[cls_id]} {conf:.2f}"

            # 绘制边界框
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            # 绘制标签背景
            (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.rectangle(image, (x1, y1 - 20), (x1 + tw, y1), color, -1)
            # 绘制标签文字
            cv2.putText(image, label, (x1, y1 - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        return image


# ========== 使用示例 ==========
if __name__ == "__main__":
    detector = YOLODetector(model_name="yolov5s", conf_thres=0.3, iou_thres=0.45)

    # 读取图像
    img = cv2.imread("test.jpg")
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # 保存临时文件供模型读取
    cv2.imwrite("temp_input.jpg", img_rgb)

    # 检测
    boxes, confs, cls_ids = detector.detect("temp_input.jpg")

    print(f"检测到 {len(boxes)} 个目标")
    for box, conf, cls_id in zip(boxes, confs, cls_ids):
        print(f"  {detector.classes[cls_id]}: {conf:.3f} at {box.tolist()}")

    # 可视化
    result_img = detector.visualize(img, boxes, confs, cls_ids)
    cv2.imwrite("output.jpg", result_img)
    print("检测结果已保存到 output.jpg")

# ========== 进阶：自定义训练 ==========
# 使用 Ultralytics YOLOv8 训练自定义数据集
from ultralytics import YOLO

# 1. 加载预训练模型 (迁移学习)
model = YOLO("yolov8n.pt")  # nano 模型

# 2. 准备数据集 (data.yaml)
# data.yaml 格式:
# path: /path/to/dataset
# train: images/train
# val: images/val
# nc: 3  # 类别数
# names: ['cat', 'dog', 'bird']

# 3. 开始训练
results = model.train(
    data="custom_data.yaml",
    epochs=100,
    imgsz=640,
    batch=16,
    device=0,        # GPU 0
    patience=20,     # 早停
    augment=True,    # 数据增强
    mosaic=0.5,      # Mosaic 增强概率
)

# 4. 验证
metrics = model.val()
print(f"mAP50: {metrics.box.map50:.4f}")
print(f"mAP50-95: {metrics.box.map:.4f}")

# 5. 导出为 ONNX (部署)
model.export(format="onnx", imgsz=640)
print("模型已导出为 yolov8n.onnx")