Pytorch 中的tensor 、ndarray、 PIL图像格式互相转换的代码

1. PIL 格式的数据和 ndarray 转换

使用PIL读取image的信息,然后转化成numpy 中的格式

image = Image.open(img_path).convert("RGB")
image = np.array(image)
## 需要将numpy 的格式转为 np.float32.这样以后转成tensor 才会是 torch.float32
image = (image / 255.0).astype(np.float32) 

2. ndarray格式的数据和 tensor转换

  1. 使用torch.from_numpy() 函数
  1. 使用transform 将ndarray 转化成Tensor 类型的变量
import albumentations as A
from albumentations.pytorch import ToTensorV2

transform = A.Compose(
        [       A.Resize(height=Image_height,width=Image_width),
            A.Normalize(
                mean=[0,0,0],
                std=[1.0,1.0,1.0],
                max_pixel_value=255.0,
            ),
            ToTensorV2(),  ## numpy 转化为tensor
        ],
    )


 augmentation = transform(image = image)
 img = augumentation['image'] ### 得到img 的tensor 类型的数据
Tensor 可以直接转为numpy 数组:
a = tensor.randn(3,3)
a.cpu().numpy()   ## 将tensor 转为numpy数组

3. Tensor格式的数据和 PIL图像数据转换

方法1:

from torchvision.transforms import ToPILImage
## image 是一个tensor 类型的数据
image = ToPILImage()(image)
image.save('1.png')

方法2:

import torchvision
 predit = torch.sigmoid(model(img))
 predit = (predit>0.5).float()
## 可以直接 调用save_image 函数保存4维度的 tensor 为image
 torchvision.utils.save_image(predit,'1.png')   

猜你喜欢

转载自blog.csdn.net/qq_41623632/article/details/126290781