[PaddleSeg 源码阅读] PaddleSeg计算 mIoU

来看下 Paddleseg/utils/metrics.py 计算 mIoU 的过程

calculate_area 函数用来计算相交的区域，

送入的 pred 和 label 形如下(1, H, W)：
(类别0到4，0为背景类)

[[[0 2 3 4 1 0 0]
  [3 2 3 4 1 4 0]
  [3 2 4 4 1 0 4]
  [3 2 3 4 1 4 1]
  [3 2 4 4 1 0 0]]]

def calculate_area(pred, label, num_classes, ignore_index=255):
    """ Calculate intersect, prediction and label area 该函数用于计算交集、prediction 和 label 的像素点数 Args: pred (Tensor): The prediction by model. label (Tensor): The ground truth of image. num_classes (int): The unique number of target classes. ignore_index (int): Specifies a target value that is ignored. Default: 255. Returns: Tensor: The intersection area of prediction and the ground on all class. Tensor: The prediction area on all class. Tensor: The ground truth area on all class """
    # 如果 pred.shape == (1, 1, 224, 224) 则将第一维去掉
    if len(pred.shape) == 4:
        pred = paddle.squeeze(pred, axis=1)
    if len(label.shape) == 4:
        label = paddle.squeeze(label, axis=1)
    if not pred.shape == label.shape:
        raise ValueError('Shape of `pred` and `label should be equal, '
                         'but there are {} and {}.'.format(pred.shape,
                                                           label.shape))
    
    # 这一步往上得到的 pred.shape == label.shape == (1, 224, 224)
    pred_area = []
    label_area = []
    intersect_area = []
    mask = label != ignore_index # 该 mask 用于将 ignore_index 忽略

    for i in range(num_classes): # 开始迭代每一个类
    
    	# 将 pred 中是第 i 类的位置拿出来
        pred_i = paddle.logical_and(pred == i, mask)
        # 将 label 中第 i 类的位置拿出来
        label_i = label == i
        # 二者做与操作，将二者相同的位置拿出来
        intersect_i = paddle.logical_and(pred_i, label_i)
        
        # bool -> int32 后求和，送入对应的list
        pred_area.append(paddle.sum(paddle.cast(pred_i, "int32")))
        label_area.append(paddle.sum(paddle.cast(label_i, "int32")))
        intersect_area.append(paddle.sum(paddle.cast(intersect_i, "int32")))
	
	# 将以上list转变为 Paddle Tensor
    pred_area = paddle.concat(pred_area)
    label_area = paddle.concat(label_area)
    intersect_area = paddle.concat(intersect_area)

    return intersect_area, pred_area, label_area

打印一下 intersect_area, pred_area, label_area 三个变量的值：

>>> intersect_area
[ 97721 147687]

>>> pred_area
[ 99482 151398]

>>> label_area
[101432 149448]

然后调用：

metrics_input = (intersect_area, pred_area, label_area)
class_iou, miou = metrics.mean_iou(*metrics_input)

def mean_iou(intersect_area, pred_area, label_area):
    """ Calculate iou. 计算 IoU Args: intersect_area (Tensor): The intersection area of prediction and ground truth on all classes. pred_area (Tensor): The prediction area on all classes. label_area (Tensor): The ground truth area on all classes. Returns: np.ndarray: iou on all classes. float: mean iou of all classes. """
    # 用 .numpy() 转化为 np.ndarray 
    intersect_area = intersect_area.numpy()
    pred_area = pred_area.numpy()
    label_area = label_area.numpy()
    
	# 计算并集
    union = pred_area + label_area - intersect_area
    
    class_iou = []
    for i in range(len(intersect_area)):
        if union[i] == 0:
        	# 如果一个该类的像素都没有，则并集
            iou = 0
        else:
        	# 计算交并比
            iou = intersect_area[i] / union[i]
        class_iou.append(iou)
    # 计算均值
    miou = np.mean(class_iou)
    return np.array(class_iou), miou