482_目标检测算法总结¶

# ---- 1. YOLOv8 推理（最快上手）----
# pip install ultralytics
from ultralytics import YOLO
from PIL import Image

# 加载预训练YOLOv8模型（n/s/m/l/x 对应从小到大）
model = YOLO('yolov8m.pt')  # m=medium，平衡速度和精度

# 对图片进行推理
results = model('image.jpg',
                conf=0.5,   # 置信度阈值（低于此的框被过滤）
                iou=0.45)   # NMS的IoU阈值

# 处理检测结果
for r in results:
    boxes = r.boxes  # 检测框对象
    for box in boxes:
        x1, y1, x2, y2 = box.xyxy[0].tolist()  # 边界框坐标
        conf = box.conf[0].item()                # 置信度
        cls = int(box.cls[0].item())             # 类别ID
        label = model.names[cls]                 # 类别名称
        print(f"{label}: {conf:.2f} @ [{x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f}]")

# 可视化（带框的图片）
annotated = results[0].plot()  # 返回带标注的numpy图像
Image.fromarray(annotated).save('result.jpg')

# ---- 2. YOLOv8 自定义数据集微调 ----
# 数据集需要YOLO格式：
# dataset/
#   images/train/*.jpg
#   images/val/*.jpg
#   labels/train/*.txt   # 每行: class_id cx cy w h (归一化坐标)
#   labels/val/*.txt

# 配置文件 data.yaml:
# train: dataset/images/train
# val: dataset/images/val
# nc: 2  # 类别数
# names: ['cat', 'dog']

model = YOLO('yolov8m.pt')  # 从预训练权重开始微调
results = model.train(
    data='data.yaml',      # 数据配置文件
    epochs=100,
    imgsz=640,             # 输入图像尺寸
    batch=16,
    lr0=0.01,              # 初始学习率
    workers=4,
    device='cuda',
    project='runs/detect', # 保存路径
    name='exp1'
)

# ---- 3. Faster RCNN（两阶段，高精度）----
import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
import torch

# 加载预训练Faster RCNN（COCO 80类）
model = fasterrcnn_resnet50_fpn(
    weights=torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights.COCO_V1
)

# 替换分类头（适配自定义类别数）
num_classes = 3  # 背景 + 2个类别（背景必须包含）
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

# 推理模式
model.eval()
import torchvision.transforms.functional as F
from PIL import Image

img = Image.open("test.jpg").convert("RGB")
img_tensor = F.to_tensor(img).unsqueeze(0)  # (1, 3, H, W)

with torch.no_grad():
    predictions = model(img_tensor)

# 解析预测结果
pred = predictions[0]
for box, label, score in zip(pred['boxes'], pred['labels'], pred['scores']):
    if score > 0.5:  # 过滤低置信度
        print(f"类别{label}: {score:.3f} @ {box.int().tolist()}")

# ---- 4. 计算 mAP 评估 ----
# pip install torchmetrics
from torchmetrics.detection.mean_ap import MeanAveragePrecision

metric = MeanAveragePrecision(iou_type="bbox")

# 格式：list of dict，每张图一个dict
preds = [{"boxes": torch.tensor([[10, 20, 100, 200]]),
           "scores": torch.tensor([0.9]),
           "labels": torch.tensor([0])}]
targets = [{"boxes": torch.tensor([[15, 25, 105, 205]]),
             "labels": torch.tensor([0])}]

metric.update(preds, targets)
result = metric.compute()
print(f"mAP@0.5: {result['map_50']:.4f}")
print(f"mAP@0.5:0.95: {result['map']:.4f}")