| import torch | |
| def pytorch_iou(pred, target, obj_num, epsilon=1e-6): | |
| ''' | |
| pred: [bs, h, w] | |
| target: [bs, h, w] | |
| obj_num: [bs] | |
| ''' | |
| bs = pred.size(0) | |
| all_iou = [] | |
| for idx in range(bs): | |
| now_pred = pred[idx].unsqueeze(0) | |
| now_target = target[idx].unsqueeze(0) | |
| now_obj_num = obj_num[idx] | |
| obj_ids = torch.arange(0, now_obj_num + 1, | |
| device=now_pred.device).int().view(-1, 1, 1) | |
| if obj_ids.size(0) == 1: # only contain background | |
| continue | |
| else: | |
| obj_ids = obj_ids[1:] | |
| now_pred = (now_pred == obj_ids).float() | |
| now_target = (now_target == obj_ids).float() | |
| intersection = (now_pred * now_target).sum((1, 2)) | |
| union = ((now_pred + now_target) > 0).float().sum((1, 2)) | |
| now_iou = (intersection + epsilon) / (union + epsilon) | |
| all_iou.append(now_iou.mean()) | |
| if len(all_iou) > 0: | |
| all_iou = torch.stack(all_iou).mean() | |
| else: | |
| all_iou = torch.ones((1), device=pred.device) | |
| return all_iou | |