| import os |
| from torchvision.io import read_image |
|
|
| class roiLeishDataset: |
| def __init__(self, file_img, file_mask, transforms=None): |
| self.file_img = file_img |
| self.file_mask = file_mask |
| self.list_img = os.listdir(self.file_img) |
| self.transforms = transforms |
|
|
| def __len__(self): |
| return len(self.list_img) |
|
|
| def __getitem__(self, idx): |
| img_path = os.path.join(self.file_img, self.list_img[idx]) |
| image = read_image(img_path) |
| mask_path = os.path.join(self.file_mask, self.list_img[idx]) |
| mask = read_image(mask_path) |
| if self.transforms: |
| image = self.transforms(image) |
| mask = self.transforms(mask) |
| mask = mask[0].unsqueeze(0) |
| return {'image':image/255.0, 'mask':mask/255.0} |
