一、YOLOv3 模型推理过程源码解析
推理过程指的是将输入图像送入训练好的YOLOv3模型，得到模型输出的预测结果。
1. 输入图像预处理 (Preprocessing)
在将图像送入模型之前，通常需要进行一系列的预处理操作，以使其符合模型的输入要求。常见的预处理步骤包括：
- 图像缩放 (Resizing): 将输入图像缩放到模型训练时所使用的尺寸，例如常见的
416x416或608x608。这通常涉及到保持图像的宽高比，并在必要时进行填充 (padding)。 - 归一化 (Normalization): 将图像的像素值归一化到
[0, 1]或[-1, 1]的范围内。这有助于加速模型收敛和提高性能。常见的归一化方法是将像素值除以255。 - 通道调整 (Channel Adjustment): 确保图像的通道顺序与模型要求的顺序一致，例如从BGR转换为RGB。
- 转换为模型输入格式: 将处理后的图像数据转换为模型所需的张量 (Tensor) 格式，并将其放置在正确的设备上 (CPU或GPU)。
源码示例 (PyTorch):
Python
import torch
import cv2
import numpy as np
def preprocess_image(image_path, input_size):
"""预处理输入图像"""
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w = img.shape[:2]
new_h, new_w = input_size
# 保持宽高比缩放
scale = min(new_h / h, new_w / w)
resized_img = cv2.resize(img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_LINEAR)
# 填充
top = (new_h - resized_img.shape[0]) // 2
bottom = new_h - resized_img.shape[0] - top
left = (new_w - resized_img.shape[1]) // 2
right = new_w - resized_img.shape[1] - left
padded_img = cv2.copyMakeBorder(resized_img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(128, 128, 128)) # 用灰色填充
# 转换为Tensor并归一化
img_tensor = torch.from_numpy(padded_img).float().permute(2, 0, 1) / 255.0
img_tensor = img_tensor.unsqueeze(0) # 添加batch维度
return img_tensor, (h, w), scale, (top, left)
# 示例用法
image_path = "test.jpg"
input_size = (416, 416)
img_tensor, original_size, scale, padding = preprocess_image(image_path, input_size)
# 将Tensor放到设备上 (如果使用GPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
img_tensor = img_tensor.to(device)
2. 模型前向传播 (Forward Pass)
预处理后的图像数据被送入加载好的YOLOv3模型中，模型会逐层进行计算，最终在不同的尺度上输出预测结果。YOLOv3通常在三个不同的尺度上进行预测，对应于网络中的三个不同的特征图。
源码示例 (PyTorch):
Python
import torch.nn as nn
# 假设已经加载了训练好的YOLOv3模型
# model = ...
# 设置模型为评估模式
model.eval()
# 关闭梯度计算，减少内存消耗并加速推理
with torch.no_grad():
outputs = model(img_tensor)
# outputs 是一个包含三个元素的列表，每个元素对应一个尺度的预测结果
# 每个元素的形状通常是 (batch_size, num_anchors_per_scale * (5 + num_classes), grid_size, grid_size)
# 例如: (1, 3 * (5 + 80), 13, 13), (1, 3 * (5 + 80), 26, 26), (1, 3 * (5 + 80), 52, 52)
3. 解析模型输出 (Output Interpretation)
模型的原始输出是三个不同尺度上的特征图。每个特征图上的每个Cell都预测了固定数量的边界框 (由anchors定义)。对于每个预测框，输出包含以下信息：
- 边界框中心坐标偏移量 (tx, ty): 相对于当前Cell的左上角。
- 边界框宽度和高度的对数偏移量 (tw, th): 相对于预定义的anchor尺寸。
- 目标置信度 (objectness score): 表示当前预测框内包含一个物体的概率。
- 类别概率 (class probabilities): 表示当前预测框内的物体属于每个类别的概率。
我们需要将这些原始输出转换为实际的边界框坐标、置信度和类别标签。
源码示例 (PyTorch):
Python
def decode_output(output, anchors, num_classes, input_size):
"""解码模型的输出"""
batch_size, _, grid_h, grid_w = output.shape
stride_h = input_size[0] / grid_h
stride_w = input_size[1] / grid_w
scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors]
output = output.view(batch_size, len(anchors), num_classes + 5, grid_h, grid_w).permute(0, 1, 3, 4, 2).contiguous()
# 中心坐标
cx = (torch.sigmoid(output[..., 0]) + torch.arange(grid_w).float().to(output.device).view(1, 1, grid_w, 1)) * stride_w
cy = (torch.sigmoid(output[..., 1]) + torch.arange(grid_h).float().to(output.device).view(1, 1, 1, grid_h).permute(0, 1, 3, 2)) * stride_h
# 宽度和高度
pw = torch.tensor([scaled_anchors[i][0] for i in range(len(anchors))]).to(output.device).view(1, len(anchors), 1, 1)
ph = torch.tensor([scaled_anchors[i][1] for i in range(len(anchors))]).to(output.device).view(1, len(anchors), 1, 1)
bw = torch.exp(output[..., 2]) * pw * stride_w
bh = torch.exp(output[..., 3]) * ph * stride_h
# 置信度
conf = torch.sigmoid(output[..., 4])
# 类别概率
pred_cls = torch.sigmoid(output[..., 5:])
return torch.cat([cx.unsqueeze(-1), cy.unsqueeze(-1), bw.unsqueeze(-1), bh.unsqueeze(-1), conf.unsqueeze(-1), pred_cls], dim=-1)
# 假设定义了anchors和num_classes
anchors = [[[116, 90], [156, 198], [373, 326]],
[[30, 61], [62, 45], [59, 119]],
[[10, 13], [16, 30], [33, 23]]]
num_classes = 80
input_size = (416, 416)
predictions = []
for i, output in enumerate(outputs):
decoded_output = decode_output(output, anchors[i], num_classes, input_size)
predictions.append(decoded_output.view(output.size(0), -1, num_classes + 5))
# 将所有尺度的预测结果合并
predictions = torch.cat(predictions, dim=1)
二、后处理模块源码解析 (Post-processing)
后处理模块的主要任务是根据模型的原始预测结果，过滤掉低置信度的预测框，并使用非极大值抑制 (NMS) 来消除冗余的重叠框，最终得到高质量的检测结果。
1. 置信度过滤 (Confidence Filtering)
首先，我们会根据预测框的目标置信度 (objectness score) 设定一个阈值，将低于该阈值的预测框过滤掉。
源码示例 (PyTorch):
Python
conf_thres = 0.5 # 置信度阈值
# 过滤掉置信度低于阈值的预测框
conf_mask = (predictions[:, :, 4] > conf_thres).unsqueeze(-1)
predictions = predictions[conf_mask.repeat(1, 1, predictions.size(-1))].view(predictions.size(0), -1, predictions.size(-1))
2. 类别概率过滤 (Class Probability Filtering)
对于剩余的预测框，我们通常会选择具有最高类别概率的类别作为该框的预测类别。然后，我们可以根据类别概率设定一个阈值，进一步过滤掉低概率的预测。
源码示例 (PyTorch):
Python
prob_thres = 0.4 # 类别概率阈值
# 获取每个预测框的最大类别概率和对应的类别索引
max_conf, max_conf_idx = torch.max(predictions[:, :, 5:], dim=-1)
max_conf = max_conf.unsqueeze(-1)
max_conf_idx = max_conf_idx.unsqueeze(-1).float()
# 将置信度和类别信息合并到预测结果中
detections = torch.cat([predictions[:, :, :4], max_conf, max_conf_idx], dim=-1)
# 过滤掉类别概率低于阈值的预测框
prob_mask = (detections[:, :, 4] > prob_thres).unsqueeze(-1)
detections = detections[prob_mask.repeat(1, 1, detections.size(-1))].view(detections.size(0), -1, detections.size(-1))
3. 非极大值抑制 (Non-Maximum Suppression, NMS)
在经过置信度和类别概率过滤后，仍然可能存在一些重叠的预测框预测到同一个物体。NMS算法的目标是选择其中置信度最高的框，并抑制掉其他与之重叠程度较高的框。
NMS算法的步骤通常如下：
- 对每个类别分别进行NMS。
- 将属于同一类别的所有预测框按照置信度从高到低排序。
- 选择置信度最高的框作为最终的检测结果，并将其加入到最终的检测列表中。
- 计算该框与其他所有框的IoU (Intersection over Union)。
- 将IoU大于某个阈值 (NMS阈值) 的框从列表中移除，因为它们与当前选择的框高度重叠，很可能是同一个物体的重复检测。
- 重复步骤3-5，直到列表为空。
源码示例 (PyTorch):
Python
def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
"""执行非极大值抑制"""
# 将预测框的格式从 (center_x, center_y, width, height) 转换为 (x1, y1, x2, y2)
box_corner = prediction.new(prediction.shape)
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
prediction[:, :, :4] = box_corner[:, :, :4]
output = [None for _ in range(len(prediction))]
for i, image_pred in enumerate(prediction):
# 过滤掉置信度低的框
conf_mask = (image_pred[:, 4] >= conf_thres).squeeze()
image_pred = image_pred[conf_mask]
if not image_pred.size(0):
continue
# 获取每个框的置信度和类别分数
class_conf, class_pred = torch.max(image_pred[:, 5:], 1, keepdim=True)
# 创建包含 (x1, y1, x2, y2, confidence, class_confidence, class_pred) 的检测结果
detections = torch.cat((image_pred[:, :5], class_conf.float(), class_pred.float()), 1)
# 获取所有唯一的类别
unique_labels = detections[:, -1].cpu().unique()
if prediction.is_cuda:
unique_labels = unique_labels.cuda()
for c in unique_labels:
# 获取属于当前类别的所有检测结果
detections_class = detections[detections[:, -1] == c]
# 按照置信度降序排序
_, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
detections_class = detections_class[conf_sort_index]
# 执行NMS
max_detections = []
while detections_class.size(0):
# 选择置信度最高的框
max_detections.append(detections_class[0].unsqueeze(0))
# 如果只有一个框了，就结束
if len(detections_class) == 1:
break
# 计算IoU
iou = bbox_iou(max_detections[-1], detections_class[1:])
# 移除IoU大于阈值的框
detections_class = detections_class[1:][iou < nms_thres]
# 将当前类别的NMS结果添加到最终输出中
if max_detections:
if output[i] is None:
output[i] = torch.cat(max_detections, dim=0)
else:
output[i] = torch.cat((output[i], torch.cat(max_detections, dim=0)), dim=0)
return output
def bbox_iou(box1, box2):
"""计算两个 bounding boxes 的 IoU"""
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
# 交集区域的坐标
inter_x1 = torch.max(b1_x1, b2_x1)
inter_y1 = torch.max(b1_y1, b2_y1)
inter_x2 = torch.min(b1_x2, b2_x2)
inter_y2 = torch.min(b1_y2, b2_y2)
# 交集区域的面积
inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)
# 并集区域的面积
b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
union_area = b1_area + b2_area - inter_area
iou = inter_area / (union_area + 1e-6) # 添加一个小的epsilon防止除零
return iou
# 应用NMS
nms_output = non_max_suppression(predictions, conf_thres=0.5, nms_thres=0.4)
# nms_output 是一个列表，每个元素对应一张输入图像的检测结果
# 每个元素是一个形状为 (num_detections, 7) 的Tensor，包含 (x1, y1, x2, y2, objectness_conf, class_conf, class_pred)
4. 后处理结果处理 (Post-NMS Processing)
经过NMS后，我们得到了最终的检测结果。这些结果通常包含边界框的坐标、目标置信度、类别置信度和类别标签。我们可能需要将这些结果转换回原始图像的尺寸，并在图像上绘制边界框和标签。
源码示例 (PyTorch):
Python
def scale_coords(coords, original_size, img_size, padding):
"""将预测框的坐标缩放回原始图像尺寸"""
h, w = original_size
img_h, img_w = img_size
top, left = padding
scale_w = w / (img_w - left - padding[1])
scale_h = h / (img_h - top - padding[0])
coords[:, [0, 2]] -= left
coords[:, [1, 3]] -= top
coords[:, [0, 2]] *= scale_w
coords[:, [1, 3]] *= scale_h
return coords
# 假设我们只有一张输入图像
if nms_output[0] is not None:
detections = nms_output[0].cpu()
# 将坐标缩放回原始图像尺寸
detections[:, :4] = scale_coords(detections[:, :4], original_size, input_size, padding)
# 在图像上绘制边界框和标签 (这部分可以使用OpenCV或其他库完成)
# ...