1

我正在使用 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 代表蒙版。

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

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

谢谢你

4

1 回答 1

4

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

您实际需要的是数据集的不同实现,每个都__getitem__返回一个图像相应的掩码。您可以在此处查看此类类的示例。

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

于 2019-02-05T06:18:11.300 回答