如何同步PyTorch与OpenCV的图像通道？

问题描述 #

使用OpenCV生成并保存图片，通过torchvision.io.read_image读取得到的张量，与直接将原numpy数组转换并permute后的张量，RGB通道存在明显差异。需要找到优雅方案，确保两种方式得到的张量完全一致，且必须保留「OpenCV创建图像、PyTorch读取」的流程。

示例代码 #

import cv2 as cv
import numpy as np
from os import getcwd
from torchvision.io import read_image
import torch as torch

pic = np.zeros((5, 6, 3))
pic[3, 0:7, 0:3] = 255
pic[1:3, 4, 0] = 255
pic[1:3, 2, 1] = 255
pic[4:6, :, 2] = 255

cwd = getcwd()
path = "/".join((cwd, "test/pic.png"))
cv.imwrite(path, pic)

read_im = read_image(path).float()

converted_im = torch.from_numpy(pic).float()
converted_im = converted_im.permute(2, 0, 1)


print("read")
print(read_im)
print()
print("converted")
print(converted_im)
print()
print("difference")
print(read_im - converted_im)

输出差异 #

read_image读取结果 #

tensor([[[  0.,   0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.,   0.,   0.],
         [255., 255., 255., 255., 255., 255.],
         [255., 255., 255., 255., 255., 255.]],
        [[  0.,   0.,   0.,   0.,   0.,   0.],
         [  0.,   0., 255.,   0.,   0.,   0.],
         [  0.,   0., 255.,   0.,   0.,   0.],
         [255., 255., 255., 255., 255., 255.],
         [  0.,   0.,   0.,   0.,   0.,   0.]],
        [[  0.,   0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0., 255.,   0.],
         [  0.,   0.,   0.,   0., 255.,   0.],
         [255., 255., 255., 255., 255., 255.],
         [  0.,   0.,   0.,   0.,   0.,   0.]]])

直接转换结果 #

tensor([[[  0.,   0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0., 255.,   0.],
         [  0.,   0.,   0.,   0., 255.,   0.],
         [255., 255., 255., 255., 255., 255.],
         [  0.,   0.,   0.,   0.,   0.,   0.]],
        [[  0.,   0.,   0.,   0.,   0.,   0.],
         [  0.,   0., 255.,   0.,   0.,   0.],
         [  0.,   0., 255.,   0.,   0.,   0.],
         [255., 255., 255., 255., 255., 255.],
         [  0.,   0.,   0.,   0.,   0.,   0.]],
        [[  0.,   0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.,   0.,   0.],
         [255., 255., 255., 255., 255., 255.],
         [255., 255., 255., 255., 255., 255.]]])

核心原因 #

1. 通道顺序不匹配：OpenCV默认以BGR通道顺序处理图像，而torchvision.io.read_image读取图像时会自动转换为RGB通道顺序的张量，导致原数组的R、B通道在读取后互换。
2. 数据类型不一致：cv.imwrite要求图像为uint8类型（PNG为8位图像），原代码中直接使用float类型数组保存，会被自动转换，但直接转torch时使用原float值，存在潜在类型差异。
3. 切片越界：原代码中0:7（超出6列）、4:6（超出5行）的切片会触发numpy的自动截断，但可能导致逻辑不符。

优雅解决方案 #

方案1：保存前统一通道顺序与数据类型 #

在保存图像前，将RGB通道转为OpenCV的BGR通道，同时转换为uint8类型，确保保存的图像与PyTorch读取后的通道顺序匹配：

# 替换原保存代码
pic_uint8 = pic.astype(np.uint8)
# 将RGB转为BGR，适配OpenCV的通道顺序
pic_bgr = cv.cvtColor(pic_uint8, cv.COLOR_RGB2BGR)
cv.imwrite(path, pic_bgr)

方案2：读取后调整通道顺序 #

如果不想修改保存流程，可在读取后将RGB转回BGR，匹配原数组的通道顺序：

# 替换原读取代码
read_im = read_image(path).float()
# 将RGB通道转为BGR，匹配原数组顺序
read_im = read_im[[2, 1, 0], :, :]

方案3：统一直接转换的类型 #

直接转换时也先转为uint8再转float，确保与保存时的数据类型一致：

# 替换原转换代码
converted_im = torch.from_numpy(pic.astype(np.uint8)).float()
converted_im = converted_im.permute(2, 0, 1)

完整修正代码 #

import cv2 as cv
import numpy as np
from os import getcwd
from torchvision.io import read_image
import torch

pic = np.zeros((5, 6, 3))
pic[3, 0:6, 0:3] = 255  # 修正切片越界
pic[1:3, 4, 0] = 255
pic[1:3, 2, 1] = 255
pic[4, :, 2] = 255  # 修正切片越界

cwd = getcwd()
path = "/".join((cwd, "test/pic.png"))

# 转换为uint8并将RGB转为BGR后保存
pic_bgr = cv.cvtColor(pic.astype(np.uint8), cv.COLOR_RGB2BGR)
cv.imwrite(path, pic_bgr)

read_im = read_image(path).float()

converted_im = torch.from_numpy(pic.astype(np.uint8)).float()
converted_im = converted_im.permute(2, 0, 1)

print("read")
print(read_im)
print()
print("converted")
print(converted_im)
print()
print("difference")
print(read_im - converted_im)

运行上述代码后，read_im与converted_im的差异将全部为0，实现完全匹配。

内容的提问来源于stack exchange，提问作者Victor Y.