根据多个条件过滤像素

Question

我有一段代码如下：

import cv2
import numpy as np
import operator
from functools import reduce

image = cv2.imread("<some image path>")
bgr   = np.int16(image)

h, w, _ = image.shape
mask    = np.zeros((h, w), np.uint8)

# Get all channels
blue  = bgr[:,:,0]
green = bgr[:,:,1]
red   = bgr[:,:,2]

rules = np.where(reduce(operator.and_, [(red > 100), (red > green), (red > blue)]

# Create mask using above rules
mask[rules] = 255

### Then use cv2.findContours ...

这段代码没有像我预期的那样运行得足够快。我想我可以通过一一应用所有条件来使其更快，即：

rule_1 = np.where(red > 100)
rule_2 = np.where(red[rule_1] > green)
rule_3 = np.where(red[rule_2] > blue)

mask[rule_3] = 255

上述方法可以加速我的代码吗？以及如何做到这一点？非常感谢！

Answer 1

一个很好的方法是（改编自 Cris Luengo 的评论）

mask = 255 * ((red > 100) & (red > green) & (red > blue)) 

但如果你需要更快，你可以使用Numba

from numba import jit, prange

@jit(nopython=True, parallel=True) 
def red_dominates(rgb, mask):
    M, N, _ = rgb.shape
    for i in prange(M):
        for j in prange(N):
            r = rgb[i,j,0]
            g = rgb[i,j,1]
            b = rgb[i,j,2]
            mask[i,j] = 255 * ((r > 100) & (r > g) & (r > b))
    return mask

请注意，使用 prange 而不是 range。这告诉 Numba 循环是可并行的。

在我的电脑上，Numba 版本的速度大约快 3 倍。

>>> bgr = np.int16(255 * np.random.random((100, 100, 3)))
>>> w = np.ones(bgr.shape[:2], np.uint8)
>>> %timeit red_dominates(bgr, w)
13.7 µs ± 26.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
>>> %timeit 255 * ((bgr[:,:,0] > 100) & (bgr[:,:,0] > bgr[:,:,1]) & (bgr[:,:,0] > bgr[:,:,2]))
46.3 µs ± 208 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

祝你好运！