用 PyTorch 处理图像

PyTorch 深度学习进阶

Michal Oleszak

Machine Learning Engineer

云数据集

云数据集示例：五张显示不同云类型的图像。

¹ https://www.kaggle.com/competitions/cloud-type-classification2/data

什么是图像？

一张云图，其中一部分被放大以显示像素。

图像由像素（"图像元素"）组成
每个像素包含颜色信息
灰度图：0–255 的整数
- 30：

灰色方块

彩色图：三个整数，分别对应红、绿、蓝通道
- RGB = (52, 171, 235)：

蓝色方块

将图像加载到 PyTorch

期望的目录结构：

clouds_train

  - cumulus

    - 75cbf18.jpg
    - ...

  - cumulonimbus
  - ...

clouds_test

  - cumulus
  - cumulonimbus
  - ...

主文件夹：clouds_train 和 clouds_test
每个主文件夹内：每类一个子文件夹
每个类别文件夹内：图像文件

将图像加载到 PyTorch

from torchvision.datasets import ImageFolder
from torchvision import transforms


train_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((128, 128)),
])


dataset_train = ImageFolder(
    "data/clouds_train",
    transform=train_transforms,
)

定义变换：
- 转为张量
- 调整为 128×128
创建数据集，传入：
- 数据路径
- 预设变换

显示图像

dataloader_train = DataLoader(
    dataset_train, 
    shuffle=True, 
    batch_size=1,
)

image, label = next(iter(dataloader_train))
print(image.shape)

torch.Size([1, 3, 128, 128])

image = image.squeeze().permute(1, 2, 0)
print(image.shape)

torch.Size([128, 128, 3])

import matplotlib.pyplot as plt
plt.imshow(image)
plt.show()

plt.show 的云图像输出

数据增强

train_transforms = transforms.Compose([

    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(45),

    transforms.ToTensor(),
    transforms.Resize((128, 128)),
])


dataset_train = ImageFolder(
    "data/clouds/train",
    transform=train_transforms,
)

数据增强：对原始图像随机变换以生成更多数据

扩大并丰富训练集
提升模型鲁棒性
降低过拟合

三张云图展示旋转变换

让我们来练习！

PyTorch 深度学习进阶