| import numpy |
| import torch |
|
|
|
|
| class AverageMeter: |
| def __init__(self, *keys): |
| self.__data = dict() |
| for k in keys: |
| self.__data[k] = [0.0, 0] |
|
|
| def add(self, dict): |
| for k, v in dict.items(): |
| self.__data[k][0] += v |
| self.__data[k][1] += 1 |
|
|
| def get(self, *keys): |
| if len(keys) == 1: |
| return self.__data[keys[0]][0] / self.__data[keys[0]][1] |
| else: |
| v_list = [self.__data[k][0] / self.__data[k][1] for k in keys] |
| return tuple(v_list) |
|
|
| def get_entire_dict_for_ddp_calculation(self): |
| return self.__data |
|
|
| def pop(self, key=None): |
| if key is None: |
| for k in self.__data.keys(): |
| self.__data[k] = [0.0, 0] |
| else: |
| v = self.get(key) |
| self.__data[key] = [0.0, 0] |
| return v |
|
|
|
|
| class ForegroundIoU(AverageMeter): |
| def __init__(self): |
| super(ForegroundIoU, self).__init__('foreground_iou') |
|
|
| def calculate_iou(self, pred, target, eps=1e-7, get_entire_list=False): |
| r""" |
| param (both hard mask): |
| pred: size [N x H x W], type: int |
| target: size [N x H x W], type: int |
| output: |
| iou: size [1] (size_average=True) or [N] (size_average=False) |
| """ |
| assert len(pred.shape) == 3 and pred.shape == target.shape, 'shape mismatch.' |
| assert pred.dtype is torch.long and target.dtype is torch.long, 'type mismatch.' |
|
|
| N = pred.size(0) |
| num_pixels = pred.size(-1) * pred.size(-2) |
| no_obj_flag = (target.sum(2).sum(1) == 0) |
|
|
| inter = (pred * target).sum(2).sum(1) |
| union = torch.max(pred, target).sum(2).sum(1) |
|
|
| inter_no_obj = ((1 - target) * (1 - pred)).sum(2).sum(1) |
| inter[no_obj_flag] = inter_no_obj[no_obj_flag] |
| union[no_obj_flag] = num_pixels |
|
|
| iou = torch.sum(inter / (union+eps)) / N |
|
|
| self.add({'foreground_iou': iou}) |
| return self.get('foreground_iou') if not get_entire_list else self.get_entire_dict_for_ddp_calculation() |
|
|
| def reset(self,): |
| super(ForegroundIoU, self).__init__('foreground_iou') |
|
|
|
|
