python——计算两个多边形的IOU

任务描述

计算两个多边形的IOU。可以是不同类的多边形，如一个矩形和一个三角形；也可以是两个同类的多边形。

代码

import cv2
import math
from skimage.draw import polygon
from skimage.feature import peak_local_max
import torch.nn.functional as F
import numpy as np

def polygon_IOU(polygon_1, polygon_2):
    """
    计算两个多边形的IOU
    :param polygon_1: [[row1, col1], [row2, col2], ...]
    :param polygon_2: 同上
    :return:
    """
    rr1, cc1 = polygon(polygon_2[:, 0], polygon_2[:, 1])
    rr2, cc2 = polygon(polygon_1[:, 0], polygon_1[:, 1])

    try:
        r_max = max(rr1.max(), rr2.max()) + 1
        c_max = max(cc1.max(), cc2.max()) + 1
    except:
        return 0

    canvas = np.zeros((r_max, c_max))
    canvas[rr1, cc1] += 1
    canvas[rr2, cc2] += 1
    union = np.sum(canvas > 0)
    if union == 0:
        return 0
    intersection = np.sum(canvas == 2)
    return intersection / union

if __name__ == '__main__':
    dx = 15
    # 三角形输入为三个点坐标
    triangle_1 = np.array([
        [200, 100],
        [180, 180],
        [220, 180]
    ])
    triangle_2 = np.array([
        [200 + dx, 100],
        [180 + dx, 180],
        [220 + dx, 180]
    ])
    # [row, col]
    rect_1 = np.array([
        [100, 100],
        [100, 200],
        [200, 200],
        [200, 100]
    ])
    rect_2 = np.array([
        [100, 150],
        [100, 250],
        [200, 250],
        [200, 150]
    ])
    iou = polygon_IOU(rect_1, rect_2)
    print('IOU={:.5f}'.format(iou))