| 114 | |
| 115 | |
| 116 | def trainMetricGPU(output, target, threshold=0.35, pr_iou=0.5): |
| 117 | assert (output.dim() in [2, 3, 4]) |
| 118 | assert output.shape == target.shape |
| 119 | output = output.flatten(1) |
| 120 | target = target.flatten(1) |
| 121 | output = torch.sigmoid(output) |
| 122 | output[output < threshold] = 0. |
| 123 | output[output >= threshold] = 1. |
| 124 | # inter & union |
| 125 | inter = (output.bool() & target.bool()).sum(dim=1) # b |
| 126 | union = (output.bool() | target.bool()).sum(dim=1) # b |
| 127 | ious = inter / (union + 1e-6) # 0 ~ 1 |
| 128 | # iou & pr@5 |
| 129 | iou = ious.mean() |
| 130 | prec = (ious > pr_iou).float().mean() |
| 131 | return 100. * iou, 100. * prec |
| 132 | |
| 133 | def ValMetricGPU(output, target, threshold=0.35): |
| 134 | assert output.size(0) == 1 |