def postprocess()

in optimum/amd/ryzenai/models/yolov5/image_processing_yolov5.py [0:0]


def postprocess(inputs, anchors, num_classes=80, stride=[8, 16, 32], shapes=[80, 40, 20]):
    nl = len(anchors)
    no = num_classes + 5

    outputs = []
    for i in range(nl):
        bs, _, ny, nx = inputs[i].shape
        grid, anchor_grid = make_grid(anchors[i], nx, ny)

        inputs[i] = inputs[i].view(bs, nl, no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

        y = inputs[i].sigmoid()
        xy = (y[..., 0:2] * 2.0 - 0.5 + grid) * stride[i]
        wh = (y[..., 2:4] * 2) ** 2 * anchor_grid
        conf = y[..., 4:]
        y = torch.cat((xy, wh, conf), -1)
        outputs.append(y.view(bs, -1, no))

    return torch.cat(outputs, 1)