深度学习模型中颜色空间转换

目录

背景

方法

结论

---

背景

情况是这样的，我训练一个深度学习模型，用于图片的处理，读取图片时使用的是skimage包，然后进行一系列常见的处理变为tensor类型送入模型中进行训练。这个模型为我自己设计的，同时还有其他的论文设计的模型，他们共用一个数据预处理，现在我不想动数据处理模块，只想在我设计的模型中进行颜色空间的转换，看看模型在另外的颜色空间中是否有不一样的表现，那么如何做呢？（如果有其他好方法或者错误记得私信我，谢谢~~！！！）

方法

所以就出现了比较蠢的方法，把即将送入模型中的tensor张量数据类型转一圈。先把batch维度删去，再从GPU中转出来，转为numpy类型，还需要通道维度的变换，最后用skimage.color包对转出来的矩阵进行颜色空间的变换。

如下，现在即将送入模型的张量为s，那么，可以在模型中使用如下代码进行颜色空间的转换，我是将ycbcr颜色空间转换为rgb颜色空间。

import skimage.color as sc
s = sr.squeeze(0)  #输入的s是即将送入模型中的用skimage读取后转换的tensor张量
s = s.cpu().numpy()  #从GPU中取出，并转换为矩阵
s = sc.ycbcr2rgb(s)*255  #颜色空间变换！255是送入模型的tensor为0-1
s = np.transpose(s,[2,0,1])  #现在的结果s就是图片RGB矩阵，值0-255
###
#你的对s的其他操作
###
s = torch.from_numpy(s).cuda() #转为tensor并放入GPU
s = s.unsqueeze(0)   #增加batch 维度，此时的s可以送入模型中训练

结论

1.注意的是cv2包读取图片时矩阵是按照BGR排列的，这和其他包不同

2.注意转换时的图片的通道、宽度、高度的排列顺序。以及浮点、整型等（在这我未考虑，对结果影响不大）。

3.注意skimage等读出图片的矩阵数值范围是0-255，而送入模型中训练的tensor值范围是0-1

4.因为我们中间将其转为了普通矩阵，因此如果是特征图，可以在此时将其可视化出来。