如何使这个循环更快？

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))

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

Answer 2

这是我根据问题陈述、Thomas Jungblut 的回答和here 的回答得出的结论。

该算法执行单像素查找并生成仅限于 COLOR_LIST 中颜色的图像。

from PIL import Image
from datetime import datetime
from math import sqrt

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)}
COLOR_CACHE = {}

def closest_color(rgb, color_list):
    if rgb not in COLOR_CACHE:
        r, g, b = rgb
        color_diffs = []
        for color in color_list:
            cr, cg, cb = color
            color_diff = sqrt(abs(r - cr)**2 + abs(g - cg)**2 + abs(b - cb)**2)
            color_diffs.append((color_diff, color))
        COLOR_CACHE[rgb] = min(color_diffs)[1]
    return COLOR_CACHE[rgb]

def clean_img_pixels(atlas_img, color_list):
    atlas_img_cleaned = atlas_img.copy()
    pixels = atlas_img_cleaned.load()
    for ii in range(atlas_img.size[0] - 1):
        for jj in range(atlas_img.size[1] - 1):
            pixel = atlas_img.getpixel((ii, jj))
            if pixel not in color_list:
                pixels[ii, jj] = closest_color(pixel, color_list)
    return atlas_img_cleaned

im = Image.open('7y1JG.png')
im = im.convert('RGB')
start_time = datetime.now()
om = clean_img_pixels(im, COLOR_LIST)
print('Time elapsed (hh:mm:ss.ms) {}'.format(datetime.now() - start_time))
om.save('7y1JG-clean.png', "PNG")

# Time elapsed (hh:mm:ss.ms) 0:00:02.932316