用python写出对yolov5s的输出端的三个featuremap的处理

以下是一个使用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)进行解码。最后，我们将解码后的信息存储到一个字典中，并将这些字典存储在一个列表中，以便后续使用。