python - 如何使这个循环更快？

Question

我希望我的图像只有 10 种特定颜色，在 color_list 中指定。
因此，我遍历每个像素，如果该像素的颜色未包含在颜色列表中，则分配相邻区域的颜色。但由于图像是 2k x 2k 像素。这个循环需要 3 分钟左右。
我确定我这样做的方式不是最佳的。我怎样才能优化我这样做的方式？

atlas_img_marked, atlas_img_cleaned = clean_img_pixels(atlas_img, color_list)

def clean_img_pixels(atlas_img, color_list):
    dd = 3
    for ii in range(atlas_img.shape[0]-1):
        for jj in range(atlas_img.shape[1]-1):
            pixelcolor = (atlas_img[ii,jj,0],atlas_img[ii,jj,1],atlas_img[ii,jj,2])
            if pixelcolor not in color_list:
                pixel2color = (atlas_img[ii-dd,jj,0],atlas_img[ii-dd,jj,1],atlas_img[ii-dd,jj,2])
                if (pixel2color == (0,0,0)) | (pixel2color not in color_list):
                    pixel2color = (atlas_img[ii+dd,jj,0],atlas_img[ii+dd,jj,1],atlas_img[ii+dd,jj,2])
                    if (pixel2color == (0,0,0)) | (pixel2color not in color_list):
                        pixel2color = (atlas_img[ii+5,jj,0],atlas_img[ii+5,jj,1],atlas_img[ii+5,jj,2])
                atlas_img_cleaned[ii,jj] = pixel2color
    return atlas_img_cleaned

更准确地说，这是耗时最长的部分:

out_colors = []
for ii in range(atlas_img.shape[0]-1):
    for jj in range(atlas_img.shape[1]-1):
        pixelcolor = (atlas_img[ii,jj,0],atlas_img[ii,jj,1],atlas_img[ii,jj,2])
        if pixelcolor not in color_list:
            out_colors.append((ii,jj))

需要 177 秒
以这种方式尝试过:

out_colors = [(ii,jj) for (ii,jj) in itertools.product(range(atlas_img.shape[0]), range(atlas_img.shape[1])) if (atlas_img[ii,jj,0],atlas_img[ii,jj,1],atlas_img[ii,jj,2]) not in color_list]

但是差别不大。需要 173 秒
这是颜色列表:

color_list = [(52, 26, 75), (9, 165, 216), (245, 34, 208), (146, 185, 85), (251, 6, 217), (223, 144, 239), (190, 224, 121), (252, 26, 157), (150, 130, 142), (51, 129, 172), (97, 85, 204), (1, 108, 233), (138, 201, 180), (210, 63, 175), (26, 138, 43), (216, 141, 61), (38, 89, 118), (0, 0, 0)]

这是一个示例图像

Answer 1

如果你放弃 numpy一起直接操作 Pillow 数组并使用元组集而不是列表，它要快得多(对我来说，这在您的示例图片上执行时间为 5 秒):

from PIL import Image
from datetime import datetime

im = Image.open('7y1JG.png')
im = im.convert('RGB')

color_list = {(52, 26, 75), (9, 165, 216), (245, 34, 208), (146, 185, 85), (251, 6, 217), (223, 144, 239),
              (190, 224, 121), (252, 26, 157), (150, 130, 142), (51, 129, 172), (97, 85, 204), (1, 108, 233),
              (138, 201, 180), (210, 63, 175), (26, 138, 43), (216, 141, 61), (38, 89, 118), (0, 0, 0)}


def clean_img_pixels(atlas_img, color_list):
    atlas_img_cleaned = atlas_img.copy().load()
    dd = 3
    for ii in range(atlas_img.size[0] - 1):
        for jj in range(atlas_img.size[1] - 1):
            if atlas_img.getpixel((ii, jj)) not in color_list:
                pixel2_color = atlas_img.getpixel((ii - dd, jj))
                if (pixel2_color == (0, 0, 0)) | (pixel2_color not in color_list):
                    pixel2_color = atlas_img.getpixel((ii + dd, jj))
                    if (pixel2_color == (0, 0, 0)) | (pixel2_color not in color_list):
                        pixel2_color = atlas_img.getpixel((ii + 5, jj))
                atlas_img_cleaned[ii, jj] = pixel2_color
    return atlas_img_cleaned


start_time = datetime.now()

out_image = clean_img_pixels(im, color_list)
time_elapsed = datetime.now() - start_time
print('Time elapsed (hh:mm:ss.ms) {}'.format(time_elapsed))

我仍然建议您进行一些额外的边界检查，因为您的图像布局方式恰好会运行。