1

我正在为分割数据集编写一个简单的自定义 DataLoader（稍后我将添加更多功能），但是我使用__getitem()__方法返回的（图像，掩码）对是不同的；返回的掩码与返回的掩码属于不同的图像。我的目录结构是/home/bohare/data/imagesand /home/bohare/data/masks。

以下是我的代码：

import torch
from torch.utils.data.dataset import Dataset
from PIL import Image
import glob
import os
import matplotlib.pyplot as plt

class CustomDataset(Dataset):
    def __init__(self, folder_path):
        
        self.img_files = glob.glob(os.path.join(folder_path,'images','*.png'))
        self.mask_files = glob.glob(os.path.join(folder_path,'masks','*.png'))
    
    def __getitem__(self, index):
        
        image = Image.open(self.img_files[index])
        mask = Image.open(self.mask_files[index])
        
        return image, mask
    
    def __len__(self):
        return len(self.img_files)

data = CustomDataset(folder_path = '/home/bohare/data')
len(data)

此代码正确给出了数据集的总大小。

但是当我使用： img, msk = data.__getitem__(n)其中 n 是任何（图像，掩码）对的索引并且我绘制图像和掩码时，它们彼此不对应。

如何修改/添加什么到代码以确保正确返回（图像，掩码）对？谢谢您的帮助。

4

1 回答 1

1

glob.glob正在无序退货，glob.glob内部调用os.listdir：

os.listdir(path) 返回一个列表，其中包含路径给定的目录中条目的名称。该列表是任意顺序的。它不包括特殊条目“。” 和 '..' 即使它们存在于目录中。

要解决它，您可以对两者进行排序，以便顺序相同：

self.img_files = sorted(glob.glob(os.path.join(folder_path,'images','*.png')))
self.mask_files = sorted(glob.glob(os.path.join(folder_path,'masks','*.png')))
于 2020-10-21T07:04:42.210 回答