神经网络输出10通道的tensor形状为(10, 512, 512),找出10通道中的最大值的通道序号,生成一个形状为(1, 512, 512)的张量,根据张量的序号存成索引图映射不同的颜色
要根据10通道中的最大值的通道序号生成一个索引图,并且将不同的通道序号映射到不同的颜色,你可以使用Python的图像处理库Pillow。首先,我们需要找出每个位置的最大值通道序号,然后使用这个序号来创建一个索引图,并映射不同的颜色。
以下是具体的实现步骤:
导入需要的库。
创建一个颜色映射字典,用于将通道序号映射到不同的颜色。
使用torch.max函数找到最大值及其索引。
使用Pillow库创建一个索引图,并应用颜色映射。
保存索引图为图片文件。
下面是具体的代码示例:
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
# 假设tensor是一个形状为(10, 512, 512)的张量
tensor = torch.randn(10, 6, 6) # 用你的张量替换这个随机生成的张量
# 创建一个颜色映射字典,你可以根据需要添加更多颜色
# 这里我们假设通道序号0到9映射到不同的颜色
color_map = {
0: (255, 0, 0), # 红色
1: (0, 255, 0), # 绿色
2: (0, 0, 255), # 蓝色
3: (255, 255, 0), # 黄色
4: (255, 0, 255), # 紫色
5: (0, 255, 255), # 青色
6: (128, 0, 0), # 棕色
7: (0, 128, 0), # 深绿色
8: (0, 0, 128), # 深蓝色
9: (128, 128, 0), # olive
}
# 使用torch.max找到最大值及其索引
_, max_channel_indices = torch.max(tensor, dim=0)
# 将索引张量转换为numpy数组,并确保数据类型为uint8
max_channel_indices_np = max_channel_indices.numpy().astype('uint8')
# 创建一个Pillow图像对象,使用L模式(灰度图像),并应用颜色映射
index_image = Image.fromarray(max_channel_indices_np, mode='L')
index_image.putpalette([v for c in color_map.values() for v in c])
plt.imshow(index_image)
plt.show()
# 保存图像
index_image.save('index_map_colored.png')
new_index_image = Image.open('index_map_colored.png')
array_image = np.array(new_index_image)
print(array_image)
查看索引图与数值一致:max_channel_indices_np
array([[5, 9, 2, 5, 0, 0],
[4, 9, 2, 7, 1, 5],
[7, 3, 9, 3, 4, 3],
[2, 8, 7, 1, 9, 6],
[3, 8, 8, 0, 5, 4],
[9, 0, 0, 0, 7, 9]], dtype=uint8)
用PS软件打开可以看到它是索引图但是有上色。
PIL重新读取后仍是索引值:
[[5, 9, 2, 5, 0, 0],
[4, 9, 2, 7, 1, 5],
[7, 3, 9, 3, 4, 3],
[2, 8, 7, 1, 9, 6],
[3, 8, 8, 0, 5, 4],
[9, 0, 0, 0, 7, 9]]