numpy - 使用 numpy 计算二进制掩码 IOU 的最快方法

Question

我正在寻找最快的方法来计算形状完全相同的 numpy 中两个二进制掩码(1 和 0 的二维数组)的并集(杰卡德索引)的交集。我的计算代码是:

import numpy as np

def binaryMaskIOU(mask1, mask2):
    mask1_area = np.count_nonzero(mask1 == 1)
    mask2_area = np.count_nonzero(mask2 == 1)
    intersection = np.count_nonzero(np.logical_and(mask1==1,  mask2==1))
    iou = intersection/(mask1_area+mask2_area-intersection)
    return iou

所以我有两个只有 1 和 0 的数组(二进制图像掩码)，并且:

• 每个掩码的面积是数组中1的数量
• 两个掩码的交集是两个掩码都等于 1 的空间位置
• IOU(交并集)定义为交集除以并集

任何优化提示/见解都值得赞赏，因为在我的代码中，我需要使用大量掩码多次运行此操作，并且我想确保它尽可能快。

Answer 1

我最初发布了一个答案，意识到我把它弄得太复杂了，当我去编辑它时发现时间比原来的函数更糟糕。已删除。

问题中的代码的性能与其他代码接近。
logic_and 中不需要两个 maskN == 1。

数据类型可以产生显着的影响。如果是 np.bool，bm1 会快得多，但这取决于对数据执行的其他操作。

如果可以更改数据类型，可能会节省大量时间，否则我看不到太多内容，也许其他人可以。

import numpy as np

def make_arrays( size = ( 5,5 ), dtype = np.int ):
    return ( np.random.randint( 2, size = size, dtype = dtype ),
             np.random.randint( 2, size = size, dtype = dtype ) )

def bm0(mask1, mask2):
    mask1_area = np.count_nonzero(mask1 == 1)       # I assume this is faster as mask1 == 1 is a bool array
    mask2_area = np.count_nonzero(mask2 == 1)
    intersection = np.count_nonzero( np.logical_and( mask1, mask2) )
    iou = intersection/(mask1_area+mask2_area-intersection)
    return iou

def bm1(mask1, mask2):
    mask1_area = np.count_nonzero( mask1 )
    mask2_area = np.count_nonzero( mask2 )
    intersection = np.count_nonzero( np.logical_and( mask1, mask2 ) )
    iou = intersection/(mask1_area+mask2_area-intersection)
    return iou

def binaryMaskIOU(mask1, mask2):   # From the question.
    mask1_area = np.count_nonzero(mask1 == 1)
    mask2_area = np.count_nonzero(mask2 == 1)
    intersection = np.count_nonzero(np.logical_and( mask1==1,  mask2==1 ))
    iou = intersection/(mask1_area+mask2_area-intersection)
    return iou


n = 50
a, b = make_arrays( ( n, n ) )

bm0( a, b )
# 0.3416313559322034

bm1(a,b)
# 0.3416313559322034

binaryMaskIOU(a,b)
# 0.3416313559322034

%timeit bm0( a, b )
#  7.73 µs ± 22.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit bm1( a, b )
# 12.2 µs ± 20.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit binaryMaskIOU( a, b )
# 10.3 µs ± 36.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

n = 500
a, b = make_arrays( ( n, n ) )

%timeit bm0( a, b )
#  342 µs ± 1.29 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit bm1( a, b )
# 1.01 ms ± 514 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit binaryMaskIOU( a, b )
#  419 µs ± 1.11 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

#################  make a and b np.bool arrays  #################
n = 500
a, b = make_arrays( ( n, n), dtype = np.bool )

%timeit bm0( a, b )
# 439 µs ± 406 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit bm1( a, b )
#  63.2 µs ± 174 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit binaryMaskIOU( a, b )
# 814 µs ± 496 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)