mirror of
https://github.com/Laurent2916/REVA-QCAV.git
synced 2024-11-14 09:08:18 +00:00
7bdac6583b
Former-commit-id: c840d14f722503d241f6bb6d899630ad6345aca0 [formerly e435a21234620add4f0e4e269a4141e5c1508cd9] Former-commit-id: 8006af185fd68cc88b2305a02513106c16758d77
33 lines
922 B
Python
33 lines
922 B
Python
import os
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset
|
|
|
|
|
|
class SphereDataset(Dataset):
|
|
def __init__(self, image_dir, transform=None):
|
|
self.image_dir = image_dir
|
|
self.transform = transform
|
|
self.images = os.listdir(image_dir)
|
|
|
|
def __len__(self):
|
|
return len(self.images)
|
|
|
|
def __getitem__(self, index):
|
|
img_path = os.path.join(self.image_dir, self.images[index])
|
|
image = np.array(Image.open(img_path).convert("RGB"), dtype=np.uint8)
|
|
|
|
mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
|
|
|
|
if self.transform is not None:
|
|
augmentations = self.transform(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
|