(prediction, num_classes, conf_thre=0.7, nms_thre=0.45)
| 31 | |
| 32 | |
| 33 | def postprocess(prediction, num_classes, conf_thre=0.7, nms_thre=0.45): |
| 34 | box_corner = prediction.new(prediction.shape) |
| 35 | box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2 |
| 36 | box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2 |
| 37 | box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2 |
| 38 | box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2 |
| 39 | prediction[:, :, :4] = box_corner[:, :, :4] |
| 40 | |
| 41 | output = [None for _ in range(len(prediction))] |
| 42 | for i, image_pred in enumerate(prediction): |
| 43 | |
| 44 | # If none are remaining => process next image |
| 45 | if not image_pred.size(0): |
| 46 | continue |
| 47 | # Get score and class with highest confidence |
| 48 | class_conf, class_pred = torch.max( |
| 49 | image_pred[:, 5 : 5 + num_classes], 1, keepdim=True |
| 50 | ) |
| 51 | |
| 52 | conf_mask = (image_pred[:, 4] * class_conf.squeeze() >= conf_thre).squeeze() |
| 53 | # _, conf_mask = torch.topk((image_pred[:, 4] * class_conf.squeeze()), 1000) |
| 54 | # Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred) |
| 55 | detections = torch.cat((image_pred[:, :5], class_conf, class_pred.float()), 1) |
| 56 | detections = detections[conf_mask] |
| 57 | if not detections.size(0): |
| 58 | continue |
| 59 | |
| 60 | nms_out_index = torchvision.ops.batched_nms( |
| 61 | detections[:, :4], |
| 62 | detections[:, 4] * detections[:, 5], |
| 63 | detections[:, 6], |
| 64 | nms_thre, |
| 65 | ) |
| 66 | detections = detections[nms_out_index] |
| 67 | if output[i] is None: |
| 68 | output[i] = detections |
| 69 | else: |
| 70 | output[i] = torch.cat((output[i], detections)) |
| 71 | |
| 72 | return output |
| 73 | |
| 74 | |
| 75 | def bboxes_iou(bboxes_a, bboxes_b, xyxy=True): |
no test coverage detected