opencv转pytorch

import torch
import cv2
img_path = "test.jpg"
#转换前: 维度顺序为h*w*3
img_cv = cv2.imread(img_path)
print("opencv: dim0={}  dim1={} dim2={}".format(img_cv.shape[0],img_cv.shape[1],img_cv.shape[2]))

#转换后： 维度顺序为3*h*w
tensor_torch = torch.FloatTensor(img_cv).permute(2, 0, 1)
print("pytorch  dim0={}  dim1={} dim2={}".format(tensor_torch.shape[0],tensor_torch.shape[1],tensor_torch.shape[2]))