优化修改PIL Image的代码：Numpy、多进程、GPU方案探讨

问题背景 #

我近期开发了一款游戏地图制作程序，游戏地图被划分为多个“Province（省份）”，由provinces.bmp文件内的同色像素组定义。程序中有三类基于PIL Image、供tkinter展示的操作存在明显性能瓶颈：

• 加载阶段：需按颜色将每个像素存入对应Province对象，当前通过遍历全部像素实现（基础游戏地图为5000×2000，共1000万次遍历）；
• 切换MapMode（地图模式）：需基于省份的文化、所有者等信息重新上色，当前需遍历每个省份的全部像素；
• 省份数据更新：用户修改省份数据后，需重新遍历该省份像素上色并更新tkinter展示。

最初考虑用Python的CUDA库，但担忧无法适配现有数据结构，且大量Province对象的像素集合难以部署，现寻求最优优化方案。

附带最小可复现示例代码：

from random import randint
import tkinter as tk
from PIL import Image
import numpy
from PIL import ImageTk
from sys import stdout

class Province:
    def __init__(self, id, colour, resource):
        self.id = id
        self.colour = colour
        self.pixels = []
        self.resource = resource

def updateProvinceOnMap(mapImage, province, resourcesToColour):
    colour = resourcesToColour[province.resource]
    for pixel in province.pixels:
        mapImage.putpixel(pixel, colour)

if __name__ == "__main__":
    colourToProvinces = dict()
    provinces = []
    provinceMapImage = Image.open("../TestMod/map/provinces.bmp")
    provinceMapArray = numpy.array(provinceMapImage)
    resourcesToColour = {"Livestock": (0, 255, 0), "Iron": (96, 96, 96), "Gold": (255, 255, 0), "Wood": (102, 51, 0)}
    resources = ["Livestock", "Iron", "Gold", "Wood"]
    resourceCount = 4
    idCounter = 1

    # Issue 1: Populating the provinces from the image
    print("Populating Provinces")
    stdout.flush()
    for y in range(0, len(provinceMapArray)):
        for x in range(0, len(provinceMapArray[y])):
            colour = (provinceMapArray[x][y][0], provinceMapArray[x][y][1], provinceMapArray[x][y][2])
            if colour not in colourToProvinces:
                resource = resources[randint(0, 3)]
                newProvince = Province(idCounter, colour, resource)
                idCounter += 1
                colourToProvinces[colour] = newProvince
                provinces.append(newProvince)
            colourToProvinces[colour].pixels.append((x, y))

    # Issue 2: Creating image from a province's field; in this case, resource
    print("Generating Map")
    stdout.flush()
    mapImage = Image.new("RGB", (len(provinceMapArray), len(provinceMapArray[0])))
    for province in provinces:
        updateProvinceOnMap(mapImage, province, resourcesToColour)
    
    window = tk.Tk()
    window.geometry("1024x776")
    canvas = tk.Canvas(window, width=len(provinceMapArray), height=len(provinceMapArray[0]))
    canvas.pack()
    canvas.create_image(0, 0, image=ImageTk.PhotoImage(mapImage))
    window.mainloop()

核心优化思路 #

用Numpy向量运算替代Python循环，将计算转移到C层执行；同时调整数据结构，避免存储大量像素坐标，改用「颜色-省份ID映射」和「ID属性数组」实现批量操作。

1. 加载阶段优化（解决1000万次遍历瓶颈） #

原代码逐个像素遍历并添加到Province的pixels列表，内存占用高且速度慢。优化后：

• 用numpy.unique一次性提取所有唯一颜色及对应像素位置
• 直接为每个颜色创建Province，无需逐个像素判断

优化代码片段：

# 替代原加载阶段的双重循环
provinceMapArray = numpy.array(provinceMapImage)
# 将RGB通道合并为单个uint32值，避免元组作为字典键的开销
flat_colours = provinceMapArray.reshape(-1, 3).view('uint32').ravel()
unique_colours, inverse_indices = numpy.unique(flat_colours, return_inverse=True)

colourToProvinces = {}
provinces = []
idCounter = 1
resources = ["Livestock", "Iron", "Gold", "Wood"]

for colour_uint32 in unique_colours:
    # 转回RGB元组
    colour = numpy.frombuffer(numpy.array(colour_uint32, dtype='uint32').tobytes(), dtype='uint8')[::-1][:3]
    colour = tuple(colour)
    resource = resources[randint(0, 3)]
    newProvince = Province(idCounter, colour, resource)
    colourToProvinces[colour_uint32] = newProvince
    provinces.append(newProvince)
    idCounter += 1

# 创建省份ID映射数组（后续上色核心依赖）
province_id_array = inverse_indices.reshape(provinceMapArray.shape[:2])

2. 切换MapMode优化（批量上色替代逐像素遍历） #

原代码遍历每个省份的所有像素调用putpixel，效率极低。优化后：

• 创建与地图同尺寸的颜色数组
• 根据省份ID对应的属性批量填充颜色
• 直接转为PIL Image，无需逐个修改像素

优化代码片段：

def generate_map_mode(province_id_array, provinces, resourcesToColour):
    # 构建ID到颜色的映射字典
    id_to_colour = {p.id: resourcesToColour[p.resource] for p in provinces}
    # 将ID数组批量转为颜色数组
    colour_array = numpy.array([id_to_colour[id] for id in province_id_array.ravel()], dtype='uint8')
    colour_array = colour_array.reshape(province_id_array.shape + (3,))
    # 直接转为PIL Image
    return Image.fromarray(colour_array)

# 生成地图模式图像
mapImage = generate_map_mode(province_id_array, provinces, resourcesToColour)

3. 省份数据更新优化（局部批量更新） #

原代码遍历单个省份的所有像素修改，优化后：

• 利用省份ID数组定位该省份的所有像素区域
• 批量修改颜色数组对应区域
• 仅更新tkinter画布的局部区域，而非全图重绘

优化代码片段：

def update_province(province_id_array, province, new_colour, mapImageArray):
    # 生成该省份的像素掩码
    mask = province_id_array == province.id
    # 批量设置颜色
    mapImageArray[mask] = new_colour
    # 计算更新区域的边界，用于tkinter局部刷新
    y_coords = numpy.where(mask.any(axis=1))[0]
    x_coords = numpy.where(mask.any(axis=0))[0]
    return (x_coords[0], y_coords[0], x_coords[-1]+1, y_coords[-1]+1)

# 使用示例：修改省份资源后更新地图
target_province = provinces[0]
target_province.resource = "Gold"
new_colour = resourcesToColour[target_province.resource]
# 假设mapImageArray是当前地图的Numpy数组
update_region = update_province(province_id_array, target_province, new_colour, mapImageArray)
# 生成局部图像并更新tkinter画布
updated_image = Image.fromarray(mapImageArray)
photo = ImageTk.PhotoImage(updated_image.crop(update_region))
canvas.create_image(update_region[0], update_region[1], image=photo, anchor=tk.NW)
canvas.image = photo  # 保留引用防止被垃圾回收

额外优化建议 #

• 避免全图重绘：每次更新仅刷新修改的局部区域，减少tkinter开销
• 内存优化：不再存储Province的pixels列表，改用ID数组映射，5000×2000地图可节省约80MB内存
• PIL批量操作：若必须使用PIL对象，改用ImageDraw.point批量绘制像素，比putpixel快10倍以上

内容的提问来源于stack exchange，提问作者Graapefruit