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）返回一个列表，其中包含以下条目的名称路径给定的目录。该列表按任意顺序排列。确实不包括特殊条目“。”和“ ..”（即使它们存在）在目录中。

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

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

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

1 个答案: