යූ-නෙට් අත්හදා බැලීම සඳහා කාර්වානා දත්ත කට්ටලය

16from torch import nn
17from pathlib import Path
18
19import torch.utils.data
20import torchvision.transforms.functional
21from PIL import Image
22
23from labml import lab

26class CarvanaDataset(torch.utils.data.Dataset):

31    def __init__(self, image_path: Path, mask_path: Path):

37        self.images = {p.stem: p for p in image_path.iterdir()}

39        self.masks = {p.stem[:-5]: p for p in mask_path.iterdir()}

42        self.ids = list(self.images.keys())

45        self.transforms = torchvision.transforms.Compose([
46            torchvision.transforms.Resize(572),
47            torchvision.transforms.ToTensor(),
48        ])

50    def __getitem__(self, idx: int):

58        id_ = self.ids[idx]

60        image = Image.open(self.images[id_])

62        image = self.transforms(image)

64        mask = Image.open(self.masks[id_])

66        mask = self.transforms(mask)

69        mask = mask / mask.max()

72        return image, mask

74    def __len__(self):

78        return len(self.ids)

82if __name__ == '__main__':
83    ds = CarvanaDataset(lab.get_data_path() / 'carvana' / 'train', lab.get_data_path() / 'carvana' / 'train_masks')