(image, mask) 对在语义分割任务中不匹配

(image, mask) pair do not match one another in a semantic segmentation task

我正在为分割数据集编写一个简单的自定义 DataLoader(稍后我将添加更多功能)但是我 return 使用 __getitem()__ 方法的(图像,掩码)对是不同的; returned 掩码与 returned 属于不同的图像。我的目录结构是 /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 是任何(图像,蒙版)对的索引,我绘制了图像和蒙版,它们彼此不对应。

我如何 modify/what 添加到代码中以确保 return 正确 return 对(图像,蒙版)?感谢您的帮助。

glob.glob 无订单退货,glob.glob calls internally os.listdir:

os.listdir(path) Return a list containing the names of the entries in the directory given by path. The list is in arbitrary order. It does not include the special entries '.' and '..' even if they are present in the directory.

为了解决这个问题,您可以对两者进行排序,以便顺序相同:

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