语义分割中的数据生成器dataloader(pytorch版)

可以参考官方文档 web documantation。主要有三个类&＃xff1a;Dataset, Sampler and DataLoader。

• Dataset:
  代表数据集的抽象类&＃xff1b;所有其他数据集都应该继承它。所有的子类都应该覆盖len&＃xff08;提供数据集的大小&＃xff09;和getitem(支持范围从0到len(self)的整形索引)。

• Sampler:
  所有采样器的基准类&＃xff1b;每个采样器子类必须提供iter方法&＃xff0c;提供一种遍历数据集元素的索引的方法&＃xff0c;以及一个返回迭代器长度的len方法。

• DataLoader:
  组合数据集和采样器&＃xff0c;并在数据集上提供单进程或多进程迭代器。

train_images_path &＃61; "./data/train_images"
train_labels_path &＃61; "./data/train_labels"class RSDataset(Dataset):def __init__(self, input_root, mode&＃61;"train", debug &＃61; False):super().__init__()self.input_root &＃61; input_rootself.mode &＃61; modeif debug &＃61;&＃61; False:self.input_ids &＃61; sorted(img for img in os.listdir(self.input_root))else:self.input_ids &＃61; sorted(img for img in os.listdir(self.input_root))[:500]self.mask_transform &＃61; transforms.Compose([transforms.Lambda(to_monochrome),transforms.Lambda(to_tensor),])self.image_transform &＃61; transforms.Compose([transforms.ToTensor(),])self.transform &＃61; DualCompose([RandomFlip(),RandomRotate90(),Rotate(),Shift(),])def __len__(self):return len(self.input_ids)def __getitem__(self, idx):# at this point all transformations are applied and we expect to work with raw tensorsimageName &＃61; os.path.join(self.input_root,self.input_ids[idx])image &＃61; np.array(cv2.imread(imageName), dtype&＃61;np.float32)mask &＃61; np.array(cv2.imread(imageName.replace("train_images", "train_labels")))/255h, w, c &＃61; image.shapemask1 &＃61; np.zeros((h, w), dtype&＃61;int)if self.mode &＃61;&＃61; "train":image, mask &＃61; self.transform(image, mask)mask1 &＃61; mask[:,:,0]return self.image_transform(image), self.mask_transform(mask1)else:mask1 &＃61; mask[:,:,0]return self.image_transform(image), self.mask_transform(mask1)###划分训练集和验证集
def build_loader(input_img_folder &＃61; "./data/train_images",batch_size &＃61; 16,num_workers &＃61; 4):# Get correct indicesnum_train &＃61; len(sorted(img for img in os.listdir(input_img_folder)))indices &＃61; list(range(num_train))seed(128381)indices &＃61; sample(indices, len(indices))split &＃61; int(np.floor(0.15 * num_train))train_idx, valid_idx &＃61; indices[split:], indices[:split]train_sampler &＃61; SubsetRandomSampler(train_idx)#set up datasetstrain_dataset &＃61; RSDataset("./data/train_images","./data/train_labels",mode &＃61; "train",)val_dataset &＃61; RSDataset("./data/train_images","./data/train_labels",mode&＃61;"valid",)valid_sampler &＃61; SubsetRandomSampler(valid_idx)train_loader &＃61; torch.utils.data.DataLoader(train_dataset, batch_size&＃61;batch_size, sampler&＃61;train_sampler,num_workers&＃61;num_workers, pin_memory&＃61;True)valid_loader &＃61; torch.utils.data.DataLoader(val_dataset, batch_size&＃61;batch_size, sampler&＃61;valid_sampler,num_workers&＃61;num_workers, pin_memory&＃61;True)return train_loader, valid_loader