这问题我太懂了——硬写if分支确实是死胡同，标签多起来根本维护不过来。咱们可以用字典来动态管理每个标签对应的画布和计数器，完美解决无限扩展的问题！

原代码的痛点 #

你现在的写法需要为每个标签手动定义画布（Zero、One）、计数器（count、count1），还要写对应的if判断。标签数量一旦增加（比如到10、100），代码会变得臃肿不堪，完全没法维护。

优化方案：用字典动态管理 #

核心思路是用两个字典分别存储每个标签对应的画布和每个标签的计数器，这样不管标签是0、1还是任意数字，都能自动适配，不用新增任何分支判断。

修改后的完整代码 #

import numpy as np
from PIL import Image
import os

# 把画布大小设为常量，方便后续修改
CANVAS_SIZE = (100, 100)
# 初始化两个字典：存储每个标签的画布、每个标签的计数
label_canvas = {}
label_counts = {}

for (x, y) in labels:
    component = uf.find(labels[(x, y)])
    labels[(x, y)] = component
    
    # 如果当前标签还没有对应的画布，就创建一个全黑画布
    if component not in label_canvas:
        label_canvas[component] = np.zeros(shape=CANVAS_SIZE, dtype=np.uint8)
        label_counts[component] = 0
    
    # 更新画布像素和计数器
    label_canvas[component][y][x] = 255
    label_counts[component] += 1
    
    # 当计数超过32时，保存图片
    if label_counts[component] > 32:
        img = Image.fromarray(label_canvas[component])
        # 动态生成文件名，比如"xxx_0.png"、"xxx_1.png"
        img.save(os.path.join(dirs, f"{file_}_{component}.png"), "png")

代码解释 #

1. 字典初始化：label_canvas用来存每个标签对应的空白画布，label_counts存每个标签的像素计数，不用提前为每个标签手动定义变量。
2. 动态创建画布：遇到新标签时，自动生成对应的全黑画布，计数器初始化为0。
3. 统一逻辑处理：不管是哪个标签，更新像素、计数、保存图片的逻辑都是统一的，完全不需要新增if分支。
4. 动态文件名：用f-string拼接文件名，直接用标签数字命名，不用手动写"Zero"、"First"这类固定名称。

额外优化建议 #

• 如果需要对某些标签做特殊处理（比如跳过某个标签），可以加个判断：if component == 特殊值: continue
• 如果保存图片后不需要继续处理该标签，可以在保存后把它从字典里移除，节省内存：del label_canvas[component], label_counts[component]

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