以下是一个简单的示例代码,用于对yolov5s输出的大小为[1,27,24,24]的featuremap进行后处理:

import numpy as np

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def post_process(output, conf_thresh=0.5, iou_thresh=0.5):
    # output: [1, 27, 24, 24]
    num_classes = 80
    num_anchors = 3
    input_size = 640
    strides = [8, 16, 32]
    anchor_sizes = [[10, 13], [16, 30], [33, 23]]
    anchor_boxes = []

    # Generate anchor boxes
    for i in range(len(strides)):
        grid_size = input_size // strides[i]
        for j in range(num_anchors):
            w, h = anchor_sizes[j]
            anchor_boxes.append([0, 0, w / strides[i], h / strides[i]])

    # Reshape output
    output = np.transpose(output, (0, 2, 3, 1))
    output = np.reshape(output, (output.shape[0], output.shape[1], output.shape[2], num_anchors, -1))

    # Apply sigmoid activation to the confidence scores
    output[..., 4] = sigmoid(output[..., 4])

    # Apply sigmoid activation to the class scores
    output[..., 5:] = sigmoid(output[..., 5:])

    # Convert the x, y coordinates to absolute values
    grid_x = np.arange(output.shape[2])
    grid_y = np.arange(output.shape[1])
    cx, cy = np.meshgrid(grid_x, grid_y)
    cx = np.reshape(cx, (-1, 1))
    cy = np.reshape(cy, (-1, 1))
    xy_grid = np.concatenate((cx, cy), axis=1)
    xy_grid = np.tile(xy_grid, (1, num_anchors))
    xy_grid = np.reshape(xy_grid, (1, output.shape[1], output.shape[2], num_anchors, 2))
    output[..., :2] = sigmoid(output[..., :2])
    output[..., :2] += xy_grid
    output[..., :2] *= strides

    # Apply exponential function to the width and height
    output[..., 2:4] = np.exp(output[..., 2:4])
    output[..., 2:4] *= anchor_boxes
    output[..., 2:4] *= strides

    # Compute the confidence scores
    conf_scores = output[..., 4]

    # Compute the class scores
    class_scores = output[..., 5:]
    class_ids = np.argmax(class_scores, axis=-1)
    class_scores = np.max(class_scores, axis=-1)

    # Filter out the detections with low confidence scores
    mask = conf_scores >= conf_thresh
    conf_scores = conf_scores[mask]
    class_ids = class_ids[mask]
    class_scores = class_scores[mask]
    boxes = output[..., :4][mask]

    # Apply non-maximum suppression to the filtered detections
    if len(boxes) > 0:
        keep = []
        x1 = boxes[..., 0]
        y1 = boxes[..., 1]
        x2 = boxes[..., 2]
        y2 = boxes[..., 3]
        area = (x2 - x1) * (y2 - y1)
        idxs = np.argsort(conf_scores)[::-1]
        while len(idxs) > 0:
            i = idxs[0]
            keep.append(i)
            xx1 = np.maximum(x1[i], x1[idxs[1:]])
            yy1 = np.maximum(y1[i], y1[idxs[1:]])
            xx2 = np.minimum(x2[i], x2[idxs[1:]])
            yy2 = np.minimum(y2[i], y2[idxs[1:]])
            w = np.maximum(0, xx2 - xx1)
            h = np.maximum(0, yy2 - yy1)
            overlap = w * h / area[idxs[1:]]
            idxs = idxs[1:][overlap <= iou_thresh]

        conf_scores = conf_scores[keep]
        class_ids = class_ids[keep]
        class_scores = class_scores[keep]
        boxes = boxes[keep]

    return boxes, class_ids, class_scores, conf_scores

该代码首先根据yolov5s的anchor box大小和步长生成anchor boxes。然后,将输出的形状转换为[1, 24, 24, 3, 85],其中85是每个anchor box的预测值数量。接下来,将置信度分数和类别分数应用于sigmoid函数,并将x、y坐标转换为绝对值。然后,将宽度和高度应用于指数函数,并将过滤掉低置信度得分的检测。最后,使用非最大抑制(NMS)将过滤后的检测进行筛选。

用python写出对yolov5s的输出的大小为1272424 的featuremap做后处理

原文地址: https://www.cveoy.top/t/topic/YhV 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录