如何让这个循环更快?

5

我希望我的图片仅有color_list中的10种特定颜色。因此,我遍历每个像素,如果该像素的颜色不包括在颜色列表中,我会将相邻区域的颜色赋值给该像素。但由于图像大小为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)]

这是一个示例图片 输入图像描述


4
color_list 转换为 set 可能会带来一些改进。 - snakecharmerb
嗯,我不确定那是否是导致它变慢的原因。 - Maryam Sadeghi
1
如果您有numpy数组,则应使用numpy函数而不是for循环。 for循环会大大减慢numpy代码的速度。 - furas
@furas 是的,图片是一个numpy数组。您可以给我一个示例,演示如何使用numpy函数代替for循环吗? - Maryam Sadeghi
1
@snakecharmerb,使用集合而不是列表确实有所不同。谢谢。使用PIL将运行时间从177秒减少到4秒。并且使用集合,进一步将其缩短到3秒。 - Maryam Sadeghi
显示剩余5条评论
2个回答

2
如果你完全放弃使用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))

我仍然建议您进行一些额外的边界检查,因为它只是因为您的图像布局方式而运行。


太棒了!谢谢Thomas。我们无法想象运行时间从177秒降至4秒! - Maryam Sadeghi

0

根据问题陈述、Thomas Jungblut的答案以及此处的答案,我得出了以下解决方案。

该算法执行单个像素查找,并生成一个仅限于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

谢谢Jeremy。但在我的情况下,周围的像素可能不一定是最接近该区域的颜色。这可能会导致许多不需要的位置出现像素颜色。 因此,在我的代码中,我尝试找到相邻的区域,并用它们替换像素颜色。但速度很慢。现在只使用PIL已经解决了这个问题。 - Maryam Sadeghi

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接