YOLOv5s 输出端特征图处理：Python 代码示例

以下是一个使用 PyTorch 实现对 YoloV5s 输出端的三个 featuremap 进行处理的示例代码：

import torch

def process_featuremaps(fmaps):
    # fmaps: 三个 featuremap 列表，每个 featuremap 的 shape 为 (batch_size, num_anchors, grid_size, grid_size, num_classes+5)
    # 5 表示 (x, y, w, h, confidence)，num_classes 表示类别数
    
    processed_fmaps = []
    
    for i, fmap in enumerate(fmaps):
        batch_size, num_anchors, grid_size, _, num_classes = fmap.shape
        
        # 将 featuremap reshape 成 (batch_size, num_anchors * grid_size * grid_size, num_classes + 5)
        fmap = fmap.view(batch_size, -1, num_classes + 5)
        
        # 取出每个 anchor 的中心点坐标 (x, y) 和宽高 (w, h)，并将它们转换为相对于图像的坐标
        x = torch.sigmoid(fmap[..., 0])
        y = torch.sigmoid(fmap[..., 1])
        w = fmap[..., 2]
        h = fmap[..., 3]
        confidence = torch.sigmoid(fmap[..., 4])
        classes = torch.sigmoid(fmap[..., 5:])
        grid_x = torch.arange(grid_size).repeat(grid_size, 1).view([1, 1, grid_size, grid_size]).float()
        grid_y = torch.arange(grid_size).repeat(grid_size, 1).t().view([1, 1, grid_size, grid_size]).float()
        anchor_w = torch.tensor([10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326]).view(1, num_anchors, 1, 1).float()
        anchor_h = torch.tensor([13, 10, 30, 16, 23, 33, 61, 30, 45, 62, 119, 59, 90, 116, 198, 156, 326, 373]).view(1, num_anchors, 1, 1).float()
        bx = (x + grid_x) / grid_size
        by = (y + grid_y) / grid_size
        bw = anchor_w * torch.exp(w)
        bh = anchor_h * torch.exp(h)
        
        # 将处理后的信息存储到一个字典中
        processed_fmap = {'x': bx, 'y': by, 'w': bw, 'h': bh, 'confidence': confidence, 'classes': classes}
        processed_fmaps.append(processed_fmap)
    
    return processed_fmaps

在该函数中，我们首先将每个 featuremap reshape 成 (batch_size, num_anchors * grid_size * grid_size, num_classes + 5) 的形状。接着，我们使用 sigmoid 函数将每个 anchor 的中心点坐标 (x, y) 和置信度 (confidence) 进行解码。然后，我们使用指数函数将每个 anchor 的宽高 (w, h) 进行解码。最后，我们将解码后的信息存储到一个字典中，并将这些字典存储在一个列表中，以便后续使用。