分类图像着色提速咨询：14400*7200大尺寸图像优化方案

我有一张尺寸为14400×7200的分类图像，使用以下代码进行着色处理时能正常运行，但耗时极长：

import numpy as np
import random
import cv2

color_map = {}
for category in np.unique(np.arange(201, dtype=np.uint8)):
    color_map[category] = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))   
    
color_map["0"]=(128, 64,128) 
color_map["1"]=(232, 35,244)
color_map["2"]=( 70, 70, 70)
color_map["3"]=(156,102,102)
color_map["4"]=(255,255,204)
color_map["5"]=(153,153,153)
color_map["6"]=(30,170, 250)
color_map["7"]=(30,170, 250)
color_map["8"]=(52,219, 163)
color_map["9"]=(0,102, 102)
color_map["10"]=( 180,130,70)
color_map["11"]=(60, 20, 220)
color_map["12"]=(60, 20, 220)
color_map["13"]=(255, 0, 0)
color_map["14"]=(255, 0, 0)
color_map["15"]=(255, 0, 0)
color_map["16"]=(255, 0, 0)
color_map["17"]=(255, 0, 0)
color_map["18"]=(255, 0, 0)
color_map["40"]=(0, 0, 0)
color_map["41"]=(0, 1, 0)
color_map["42"]=(0, 2, 0)
color_map["43"]=(0, 3, 0)
color_map["44"]=(0, 4, 0)
color_map["45"]=(0, 5, 0)
color_map["46"]=(0, 6, 0)
color_map["47"]=(0, 7, 0)
color_map["48"]=(0, 8, 0)
color_map["49"]=(0, 9, 0)
color_map["50"]=(0, 10, 0)
color_map["51"]=(0, 11, 0)

a_array = np.array(a)
category_image = a_array
colored_image = np.zeros((category_image.shape[0], category_image.shape[1], 3), dtype=np.uint8)
for h in range(category_image.shape[0]):
    for g in range(category_image.shape[1]):
        category = category_image[h, g]
        colored_image[h, g] = color_map[str(category)]
cv2.imwrite(j, colored_image)

核心优化思路 #

原代码的性能瓶颈主要是Python层面的逐像素双重循环（处理超1亿像素），以及字典键类型不统一导致的额外类型转换。以下是针对性优化方案：

1. 构建颜色查找表（LUT）替代字典 #

利用分类标签是uint8类型（取值0-255）的特性，直接创建一个长度为256的数组作为颜色查找表，索引对应分类标签，值为RGB颜色。相比字典，数组索引的速度快几个数量级。

2. 统一键类型，避免字符串转换 #

原代码中颜色字典同时存在数字键和字符串键，每次取值都要把标签转为字符串，完全可以统一用数字键，省去类型转换开销。

3. 用numpy向量化操作替代逐像素循环 #

numpy的向量化操作是C语言实现的，能一次性处理整个数组，彻底避免Python循环的低效问题。

优化后的完整代码 #

import numpy as np
import random
import cv2

# 1. 初始化颜色查找表（长度256，对应uint8的所有可能值）
color_lut = np.zeros((256, 3), dtype=np.uint8)

# 2. 为0-200的分类随机生成颜色
for category in range(201):
    color_lut[category] = [random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)]

# 3. 覆盖指定分类的颜色（直接用数字索引，无需字符串键）
color_lut[0] = (128, 64, 128)
color_lut[1] = (232, 35, 244)
color_lut[2] = (70, 70, 70)
color_lut[3] = (156, 102, 102)
color_lut[4] = (255, 255, 204)
color_lut[5] = (153, 153, 153)
color_lut[6] = (30, 170, 250)
color_lut[7] = (30, 170, 250)
color_lut[8] = (52, 219, 163)
color_lut[9] = (0, 102, 102)
color_lut[10] = (180, 130, 70)
color_lut[11] = (60, 20, 220)
color_lut[12] = (60, 20, 220)
color_lut[13] = (255, 0, 0)
color_lut[14] = (255, 0, 0)
color_lut[15] = (255, 0, 0)
color_lut[16] = (255, 0, 0)
color_lut[17] = (255, 0, 0)
color_lut[18] = (255, 0, 0)
color_lut[40] = (0, 0, 0)
# 处理41-51的渐变绿色，简化循环
for i in range(41, 52):
    color_lut[i] = (0, i - 40, 0)

# 4. 加载分类图像（显式指定uint8类型，减少内存占用）
category_image = np.array(a, dtype=np.uint8)

# 5. 向量化着色：直接通过索引查找颜色，无需逐像素循环
colored_image = color_lut[category_image]

# 6. 保存图像
cv2.imwrite(j, colored_image)

额外优化建议 #

• 确保输入的category_image是uint8类型，减少不必要的内存占用和类型转换。
• 若图像尺寸过大导致内存紧张，可以分块处理图像，每次处理一部分后再拼接结果。

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