Python代码优化：使用数字替换像素值，简化图像处理

下面是修改后的代码，将保存文件中的像素值替换为对应的数字：

from PIL import Image
import os

# 输入文件夹和输出文件夹的路径
input_folder = r'C:\Users\jh\Desktop\data\images_block'
output_folder = r'C:\Users\jh\Desktop\data\images_flatten'

# 确保输出文件夹存在
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 定义要提取的像素值和对应的数字映射关系
pixels_to_extract = {(181, 0, 0): 4, (232, 14, 14): 3, (255, 208, 69): 2, (79, 210, 125): 1}

# 循环遍历输入文件夹中的所有PNG图片
for filename in os.listdir(input_folder):
    if filename.endswith('.png'):
        input_path = os.path.join(input_folder, filename)
        
        # 打开PNG图片并将其转换为RGB模式
        image = Image.open(input_path).convert('RGB')
        
        # 获取图片的像素值
        pixels = list(image.getdata())
        
        # 创建空字典用于统计四类像素的数量
        pixel_counts = {pixel: 0 for pixel in pixels_to_extract}
        
        # 统计四类像素的数量
        for pixel in pixels:
            if pixel in pixel_counts:
                pixel_counts[pixel] += 1
        
        # 提取指定的像素值并替换其他像素为(0, 0, 0)
        extracted_pixels = [(pixels_to_extract.get(pixel, 0), 0, 0) for pixel in pixels]
        
        # 将像素值保存到输出文件夹中的文本文件中
        output_path = os.path.join(output_folder, filename.replace('.png', '.txt'))
        with open(output_path, 'w') as file:
            for pixel in extracted_pixels:
                file.write(f'{pixel[0]}, {pixel[1]}, {pixel[2]}
')
        
        # 输出当前图片的四类像素的数量
        print(f'{filename}的像素数量统计结果：')
        for pixel, count in pixel_counts.items():
            print(f'{pixel}: {count}')
        print('')

print('转换完成')

在这个修改后的代码中，通过使用pixels_to_extract.get(pixel, 0)来替换像素值。如果像素值不在pixels_to_extract字典中，则返回0。实际上，这样可以避免KeyError异常。