PyTorch 在加载图像/掩码文件以进行图像分割时如何处理标签？

Posted 2023-03-27

【中文标题】PyTorch 在加载图像/掩码文件以进行图像分割时如何处理标签？

【英文标题】：How does PyTorch handle labels when loading image/mask files for image segmentation?

【发布时间】：2019-06-28 21:55:32

【问题描述】：

我正在使用 PyTorch 开始一个图像分割项目。我在一个文件夹和 2 个子文件夹中有一个缩减的数据集 - “图像”用于存储图像，“掩码”用于蒙版图像。图像和蒙版是具有 3 个通道和 256x256 像素的 .png 文件。因为它是图像分割，所以必须逐个像素地进行标记。为简单起见，我目前只使用 2 个类。到目前为止，我实现了以下目标：

我能够通过

将我的文件加载到“图像”或“掩码”类中

root_dir="./images_masks"
train_ds_untransf = torchvision.datasets.ImageFolder(root=root_dir)
train_ds_untransf.classes
Out[621]:
['images', 'masks']  

并将数据转换为张量

from torchvision import transforms
train_trans = transforms.Compose([transforms.ToTensor()])
train_dataset = torchvision.datasets.ImageFolder(root=root_dir,transform=train_trans)

这个“train_dataset”中的每个张量都具有以下形状：

train_dataset[1][0].shape
torch.Size([3, 256, 256])

现在我需要将加载的数据输入 CNN 模型，并为此探索了 PyTorch DataLoader

train_dataloaded = DataLoader(train_dataset, batch_size=2, shuffle=False, num_workers=4)

我使用以下代码检查生成的张量的形状

for x, y in train_dl:
    print (x.shape)
    print (y.shape)
    print(y)

得到

torch.Size([2, 3, 256, 256])
torch.Size([2])
tensor([0, 0])
torch.Size([2, 3, 256, 256])
torch.Size([2])
tensor([0, 1])
.
.
.

形状似乎正确。但是，第一个问题是我得到了同一个文件夹的张量，由一些具有相同值 [0, 0] 的“y”张量表示。我希望它们都是 [1, 0]：1 代表图像，0 代表蒙版。

第二个问题是，虽然当标签是整个图像时文档很清楚，但不清楚如何将其应用于像素级别的标签，我确信标签不正确。

正确标记此数据集的替代方法是什么？

谢谢

【参考方案1】：

torchvision.datasets.ImageFolder 类是为图像分类 问题设计的，而不是用于分割的；因此，它期望每个图像有一个整数标签，并且标签由存储图像的子文件夹确定。因此，就您的数据加载器而言，您有两类图像“图像”和“掩码”，您的网络会尝试区分它们。

您真正需要的是一个不同的数据集实现，它为每个__getitem__ 返回一个图像和相应的掩码。您可以查看此类类的示例here。

此外，您的二进制像素级标签存储为 3 通道图像有点奇怪。分割掩码通常存储为单通道图像。

【讨论】：

分割掩码以 3 通道图像的形式提供是很常见的，其中所有通道在像素方面都是相等的。我相信 .png 可以原生存储单通道图像，但有些人就是不知道，或者出于任何原因使用 3 个通道。我不会担心 OPs 的位置。

@Jatentaki 分割掩码在存储为图像时，通常存储为 索引 RGB（PIL Image class 中的'P' 模式）。因此，尽管每个像素可以获得 3 个通道，但它实际上只有通道

@Shai 您能否更具体地向我指出您链接的 GitHub 项目中正确的类示例？非常感谢！

@Fabio 该存储库中的不同 .py 文件是为不同语义分割基准量身定制的 Dataset 实例。