import os
import cv2
import numpy as np
import matplotlib.pyplot as plt

# 输入文件夹路径
input_folder = r'D:\Users\Desktop\img'
# 输出文件夹路径
output_folder = r'D:\Users\Desktop\output_images'

# 如果输出文件夹不存在，则创建它
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 遍历输入文件夹中的所有文件
for filename in os.listdir(input_folder):
    if filename.endswith(('.jpg', '.png', '.jpeg')):
        # 构建完整的文件路径
        input_path = os.path.join(input_folder, filename)
        # 读取图像
        im2_uint8 = cv2.imread(input_path, cv2.IMREAD_GRAYSCALE)

        # 创建自定义 LUT
        lut = np.zeros((256, 1, 3), dtype=np.uint8)
        for i in range(256):
            if i < 32:
                # 黑到红渐变
                lut[i, 0, 2] = 255  # 红色通道
                lut[i, 0, 0] = 0    # 蓝色通道
                lut[i, 0, 1] = 0    # 绿色通道
            else:
                # 红到黄渐变（红+绿=黄）
                lut[i, 0, 2] = 255      # 红色通道保持最大
                lut[i, 0, 1] = (i - 128) * 2    # 绿色通道逐渐增加

        # 转换为 BGR 并应用 LUT
        pseudo_color1 = cv2.LUT(cv2.cvtColor(im2_uint8, cv2.COLOR_GRAY2BGR), lut)

        # 构建输出文件路径
        output_filename = os.path.splitext(filename)[0] + '_pseudo_color.jpg'
        output_path = os.path.join(output_folder, output_filename)

        # 保存处理后的图像
    #    cv2.imwrite(output_path, pseudo_color1)

        # 使用 matplotlib 显示图像（可选）
        #plt.imshow(pseudo_color1)
        #plt.title('Pseudo Color Image')
        #plt.axis('off')  # 不显示坐标轴
        #plt.show()
        # 提取原始文件名（不包含扩展名）
        base_name = os.path.splitext(filename)[0]
        # 构建输出文件路径
        output_filename = f'{base_name}_pseudo_color.png'
        output_path = os.path.join(output_folder, output_filename)

        # 使用 matplotlib 显示并保存图像
        plt.imshow(pseudo_color1, cmap='viridis')
        plt.gca().set_axis_off()
        plt.savefig(output_path, bbox_inches='tight', pad_inches = 0)
        plt.close()  # 关闭当前图形，避免图形重叠


#0 COLORMAP_AUTUMN
#1 COLORMAP_BONE
#2 COLORMAP_JET
#3 COLORMAP_WINTER
#4 COLORMAP_RAINBOW
#5 COLORMAP_OCEAN
# COLORMAP_SUMMER
#7 COLORMAP_SPRING
#8 COLORMAP_COOL
#9 COLORMAP_HSV
#10 COLORMAP_PINK
#11 COLORMAP_HOT
#pseudo_color = cv2.applyColorMap(im2_uint8, cv2.COLORMAP_HSV)
#COLORMAP_JET：蓝-青-黄-红渐变（类似热力图）
#COLORMAP_HOT：黑-红-黄-白渐变
#COLORMAP_HSV：HSV色彩空间循环
#其他选项如COLORMAP_BONE, COLORMAP_COOL, COLORMAP_PINK等。