mirror of
https://github.com/Laurent2916/REVA-QCAV.git
synced 2024-11-14 00:58:15 +00:00
5f46efa5a1
Former-commit-id: 7b293e392cc7d4135ef8562faece6f491c623718 [formerly 381a418ceab2cb7f367f07b9f3ea4f3c6a41ecac] Former-commit-id: c2ddd57c4a3592c93640170a737506bb64b60864
46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
from pathlib import Path
|
|
|
|
import albumentations as A
|
|
import numpy as np
|
|
from albumentations.pytorch import ToTensorV2
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset
|
|
|
|
|
|
class SphereDataset(Dataset):
|
|
def __init__(self, image_dir, transform=None):
|
|
self.images = list(Path(image_dir).glob("**/*.jpg"))
|
|
self.transform = transform
|
|
|
|
def __len__(self):
|
|
return len(self.images)
|
|
|
|
def __getitem__(self, index):
|
|
image = np.array(Image.open(self.images[index]).convert("RGB"), dtype=np.uint8)
|
|
|
|
if self.transform is not None:
|
|
mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
|
|
augmentations = self.transform(image=image, mask=mask)
|
|
image = augmentations["image"]
|
|
mask = augmentations["mask"]
|
|
else:
|
|
mask_path = self.images[index].parent.joinpath("MASK.PNG")
|
|
mask = np.array(Image.open(mask_path).convert("L"), dtype=np.uint8) / 255
|
|
|
|
preprocess = A.Compose(
|
|
[
|
|
A.SmallestMaxSize(1024),
|
|
A.ToFloat(max_value=255),
|
|
ToTensorV2(),
|
|
],
|
|
)
|
|
augmentations = preprocess(image=image, mask=mask)
|
|
image = augmentations["image"]
|
|
mask = augmentations["mask"]
|
|
|
|
# make sure image and mask are floats
|
|
image = image.float()
|
|
mask = mask.float()
|
|
|
|
return image, mask
|