深度学习（np,tensor,pil,cv2互转）

Dsp Tian / 2024-10-25 / 原文

在做深度学习的时候经常会在这四种格式上互转。

转换时有几个需要注意的点：

1. np和cv2都能通过cv2.imshow显示出来。

2. torchvision的transforms会把HWC转为CHW。

3. from_numpy读取的HWC还是HWC，需要额外对tensor旋转。

4. pil和cv2读取的都是HWC格式，不过C中的BGR和RGB顺序不同。

下面是两两互转的例子：

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

img_pil = Image.open('1.jpg')
img_cv2 = cv2.imread('1.jpg')

#pil <-> np
img_np = np.array(img_pil)
img_pil = Image.fromarray(img_np)

#pil <-> torch
img_tensor = torchvision.transforms.ToTensor()(img_pil)
img_pil = torchvision.transforms.ToPILImage()(img_tensor)

#pil <-> cv2
img_cv2 = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
img_pil = Image.fromarray(cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB))

#cv2 <-> np
img_np = np.array(img_cv2)
img_cv2 = img_np                                                #直接给能imshow出来
# img_cv2 = cv2.cvtColor(np.array(img_cv2),cv2.COLOR_BGR2RGB)   #网上常见的方法转一次颜色不对
# img_cv2 = cv2.cvtColor(img_cv2,cv2.COLOR_BGR2RGB)             #需要再转一次  

# cv2 <-> torch
img_tensor = torchvision.transforms.ToTensor()(img_cv2)
img_cv2 = img_tensor.permute(1,2,0).numpy()

# np <-> torch
img_tensor = torch.from_numpy(img_np)
img_np = img_tensor.numpy()

# img_pil.show()
cv2.imshow('img',img_cv2)
cv2.waitKey(0)