减少嵌套 for 循环的执行时间

时间:2021-01-11 12:05:57

标签: python numpy image-processing

目前,下面提到的代码需要大约 6.5 秒才能执行。 (1080, 1920) 中 image_mask 的形状。

pixel_array = np.zeros(image_mask.shape, dtype=np.uint8)
for i in range(image_mask.shape[0]):
    for j in range(image_mask.shape[1]):
        if (image_mask[i, j, :] == 0).all():
            pixel_array[i, j, :] = 255
        else:
        # elif (image_mask[i, j, :] == 255).all():
            pixel_array[i, j, :] = 0

是否有一种更快的方法可以在 0.5 秒或更短的时间内获得处理后的图像?

1 个答案:

答案 0 :(得分:0)

您可以使用布尔掩码来做到这一点。

import numpy as np

image_mask = np.random.randint(0, 255, size=(100, 200, 3))
image_mask[0, 0, :] = 0
bool_mask = np.all(image_mask==0, axis=-1)
image_mask[bool_mask] = 255
image_mask[~bool_mask] = 0

比较 -

import numpy as np

def f1(img):
    bool_mask = np.all(img==0, axis=-1)
    img[bool_mask] = 255
    img[~bool_mask] = 0
    return img

def f2(img):
    pixel_array = np.zeros(img.shape, dtype=np.uint8)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            if (img[i, j, :] == 0).all():
                pixel_array[i, j, :] = 255
            else:
            # elif (image_mask[i, j, :] == 255).all():
                pixel_array[i, j, :] = 0
    return pixel_array

image_mask = np.random.randint(0, 255, size=(1080, 1920, 3))
image_mask[0, 0, :] = 0

np.allclose(f1(image_mask.copy()), f2(image_mask))

%timeit f1(image_mask)
%timeit f2(image_mask)

给予 -

66 ms ± 572 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
5.61 s ± 45.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)