plt.imsave与OpenCV图像转换差异及无IO等效预处理方案问询

背景与代码 #

原输入的spectrogram是形状为(2,X)、数值范围在-80到+50之间的双通道张量。以下是两段处理代码：

代码1（训练时使用） #

plt.imsave(f"{save_path}/img.png", spectrogram[0]) # 保存为4通道RGBA图像
spectrogram1 = torchvision.io.read_image(f"{save_path}/img.png", io.ImageReadMode.GRAY)[0].float()
spectrogram1 = torch.stack([spectrogram1, spectrogram1])

处理后spectrogram1的数值范围为0-255，可视化效果与原图像一致。

代码2（推理时尝试替代） #

spectrogram2 = cv2.cvtColor(spectrogram[0], cv2.COLOR_GRAY2RGBA)
spectrogram2 = cv2.cvtColor(spectrogram2, cv2.COLOR_RGBA2GRAY)
spectrogram2 = torch.stack([spectrogram2, spectrogram2])

spectrogram2数值范围同样为0-255，视觉效果与spectrogram1高度相似，但基于spectrogram1训练的分类模型对其输出随机结果。

问题 #

1. 为何spectrogram1与spectrogram2存在差异？
2. 推理阶段如何实现与spectrogram1完全一致的转换，且避免文件IO带来的开销？

解答 #

1. 两者差异的核心原因 #

两段代码的处理逻辑完全不同，导致最终张量的数值分布存在本质区别：

• 代码1的处理流程：
  1. plt.imsave接收单通道数值数组时，默认会应用**默认颜色映射（Colormap，如viridis）**将数值转换为彩色RGBA图像。这个过程会把原-80到50的数值范围线性映射到0-255，再通过颜色映射生成RGB通道的彩色像素值。
  2. torchvision.io.read_image以GRAY模式读取时，会将彩色图像按照标准灰度转换公式（Y = 0.299*R + 0.587*G + 0.114*B）转换为单通道灰度图，忽略Alpha通道。
• 代码2的处理流程：
  1. cv2.COLOR_GRAY2RGBA只是将单通道数组复制到RGBA的四个通道中，没有应用任何颜色映射。
  2. cv2.COLOR_RGBA2GRAY转换时，直接取原通道的数值进行简单加权（或直接复用单通道值，取决于输入类型），完全没有经过颜色映射的步骤。

尽管两者的数值范围都是0-255且视觉相似，但内部每个位置的具体数值完全不同，模型训练时学习到的是代码1生成的特征分布，自然无法识别代码2的输入。

2. 无IO开销的等效转换实现 #

要复现代码1的效果，需要在内存中完整模拟plt.imsave的颜色映射+灰度转换流程，无需写入文件。具体实现代码如下：

import matplotlib.pyplot as plt
import torch

def convert_spectrogram(spec):
    # 1. 获取matplotlib默认颜色映射（与plt.imsave一致）
    cmap = plt.get_cmap()
    # 2. 将输入数值范围[-80,50]归一化到[0,1]
    norm_spec = (spec - (-80)) / (50 - (-80))
    # 3. 应用颜色映射，得到RGBA格式的numpy数组（形状为(X, 4)，数值0-1）
    rgba = cmap(norm_spec)
    # 4. 按照torchvision的灰度转换公式，将RGBA转为灰度（忽略Alpha通道）
    gray = 0.299 * rgba[..., 0] + 0.587 * rgba[..., 1] + 0.114 * rgba[..., 2]
    # 5. 转换为0-255的float张量
    gray_tensor = torch.tensor(gray * 255, dtype=torch.float32)
    # 6. 堆叠为双通道
    return torch.stack([gray_tensor, gray_tensor])

# 使用示例
spectrogram1_equivalent = convert_spectrogram(spectrogram[0])

关键细节说明：

• 必须使用与plt.imsave一致的颜色映射（默认是viridis，若训练时修改过cmap参数，需同步调整）。
• 归一化步骤必须严格对应原数值范围[-80,50]，确保与plt.imsave的自动归一化逻辑一致。
• 灰度转换公式必须与torchvision.io.read_image的GRAY模式对齐，避免因转换规则差异导致数值偏差。

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