（图像，掩模）对在语义分割任务中彼此不匹配

Question

我正在为分割数据集编写一个简单的自定义 DataLoader（稍后我将添加更多功能），但我使用__getitem()__方法返回的（图像，掩码）对是不同的；返回的掩码属于与返回的图像不同的图像。我的目录结构是/home/bohare/data/images和/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 是任何（图像，掩模）对的索引并且我绘制图像和掩模时，它们彼此不对应。

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

Answer 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')))