如何在Pythorch中为图像及其遮罩定制数据集?

2021-05-13 14:51:36 发布

您现在位置:Python中文网/ 问答频道 /正文

我有两个tif图像的数据集文件夹,一个是名为BMMCdata的文件夹,另一个是BMMCdata图像的掩码BMMCmasks(图像的名称是对应的)。我试图制作一个定制的数据集,并随机分割数据进行训练和测试。我现在犯了个错误

self.filenames.append(fn)
AttributeError: 'Customdataset' object has no attribute 'filenames'

任何评论都将不胜感激。在

^{pr2}$
1条回答
网友
1楼 ·

Pythorch社区的@ptrblck给出了答案。谢谢你

 # get all the image and mask path and number of images
 folder_data = glob.glob("D:\\Neda\\Pytorch\\U-net\\BMMCdata\\data\\*.tif")
 folder_mask = glob.glob("D:\\Neda\\Pytorch\\U-net\\BMMCmasks\\masks\\*.tif")

 # split these path using a certain percentage
 len_data = len(folder_data)
 print(len_data)
 train_size = 0.6

 train_image_paths = folder_data[:int(len_data*train_size)]
 test_image_paths = folder_data[int(len_data*train_size):]

 train_mask_paths = folder_mask[:int(len_data*train_size)]
 test_mask_paths = folder_mask[int(len_data*train_size):]


 class CustomDataset(Dataset):
    def __init__(self, image_paths, target_paths, train=True):   # initial logic 
      happens like transform

         self.image_paths = image_paths
         self.target_paths = target_paths
         self.transforms = transforms.ToTensor()

    def __getitem__(self, index):

        image = Image.open(self.image_paths[index])
        mask = Image.open(self.target_paths[index])
        t_image = self.transforms(image)
     return t_image, mask

def __len__(self):  # return count of sample we have

    return len(self.image_paths)

train_dataset = CustomDataset(train_image_paths, train_mask_paths, train=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=1)

test_dataset = CustomDataset(test_image_paths, test_mask_paths, train=False)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=4, shuffle=False, num_workers=1)

相关问题