Pytorch：將影象tensor資料用Opencv顯示

Pytorch：將影象tensor資料用Opencv顯示 - 知乎 (zhihu.com)

將影象tensor資料用Opencv顯示

首先匯入相關庫：*

import torch
from torchvision import transforms
from PIL import Image
import numpy as np
import cv2

利用PIL中的Image開啟一張圖片

image2=Image.open('pikachu.jpg')

這裡print看一下image2的影象資料型別，這裡可以直接呼叫image2.show()直接顯示：

print(image2)
<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=814x982 at 0x1E988BA74A8>

將image2轉化為tensor資料（為什麼轉化為tensor，當然是為了方便計算）

transform2=transforms.Compose([transforms.ToTensor()])
tensor2=transform2(image2)
print('tensor2:',tensor2)#列印看一下tensor的資料
print(tensor2.dtype)#torch.float32
print(tensor2.shape)#返回tensor2_shape torch.Size([3, 982, 814])->3通道982*814的RGB影象
tensor2: tensor([[[0.5647, 0.5686, 0.5686,  ..., 0.5725, 0.5725, 0.5686],
         [0.5647, 0.5686, 0.5725,  ..., 0.5725, 0.5725, 0.5686],
         [0.5686, 0.5725, 0.5725,  ..., 0.5725, 0.5725, 0.5686],
         ...,
         [0.3176, 0.3216, 0.3216,  ..., 0.3098, 0.3098, 0.3098],
         [0.3176, 0.3216, 0.3216,  ..., 0.3098, 0.3098, 0.3098],
         [0.3176, 0.3216, 0.3216,  ..., 0.3137, 0.3137, 0.3137]],

        [[0.7412, 0.7451, 0.7451,  ..., 0.7490, 0.7490, 0.7451],
         [0.7412, 0.7451, 0.7490,  ..., 0.7490, 0.7490, 0.7451],
         [0.7451, 0.7490, 0.7490,  ..., 0.7490, 0.7490, 0.7451],
         ...,
         [0.5529, 0.5569, 0.5569,  ..., 0.5451, 0.5451, 0.5451],
         [0.5529, 0.5569, 0.5569,  ..., 0.5451, 0.5451, 0.5451],
         [0.5529, 0.5569, 0.5569,  ..., 0.5490, 0.5490, 0.5490]],

        [[0.9059, 0.9098, 0.9098,  ..., 0.9098, 0.9098, 0.9059],
         [0.9059, 0.9098, 0.9137,  ..., 0.9098, 0.9098, 0.9059],
         [0.9098, 0.9137, 0.9137,  ..., 0.9098, 0.9098, 0.9059],
         ...,
         [0.8275, 0.8314, 0.8314,  ..., 0.8275, 0.8275, 0.8275],
         [0.8275, 0.8314, 0.8314,  ..., 0.8275, 0.8275, 0.8275],
         [0.8275, 0.8314, 0.8314,  ..., 0.8314, 0.8314, 0.8314]]])
 

要將tensor影象資料轉為opencv支援的影象資料，首先要了解opencv所支援的影象資料：

image3=cv2.imread('pokeman/pikachu/00000000.jpg')
print(image3)
[[[231 189 144]
  [232 190 145]
  [232 190 145]
  ...
  [232 191 146]
  [232 191 146]
  [231 190 145]]]
print(image3.shape)
(982, 814, 3)
print(type(image3))
<class 'numpy.ndarray'>
print(image3.dtype)
uint8
 

所以我們知道opencv支援的影象資料時numpy格式，資料型別為uint8，而且畫素值分佈在[0,255]之間。 但是從上面的tensor資料可以看出，畫素值並不是分佈在[0,255]，且資料型別為float32,所以需要做一下normalize和資料變換，將影象資料擴充套件到[0,255]。還有一點不同的是tensor(3,982, 814)、numpy(982, 814, 3)儲存的資料維度順序不同。

array1=tensor2.numpy()#將tensor資料轉為numpy資料
maxValue=array1.max()
array1=array1*255/maxValue#normalize，將影象資料擴充套件到[0,255]
mat=np.uint8(array1)#float32-->uint8
print('mat_shape:',mat.shape)#mat_shape: (3, 982, 814)
mat=mat.transpose(1,2,0)#mat_shape: (982, 814，3)
cv2.imshow("img",mat)
cv2.waitKey()

這是由於opencv中的顏色通道順序是BGR而PIL、torch裡面的影象顏色通道是RGB，利用cvtColor對顏色通道進行轉換

mat=cv2.cvtColor(mat,cv2.COLOR_BGR2RGB)
cv2.imshow("img",mat)
cv2.waitKey()