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

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 calls internally os.listdir:

os.listdir(path) 返回包含条目名称的列表 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')))