Laurent/REVA-QCAV
1
0
Fork 0
mirror of https://github.com/Laurent2916/REVA-QCAV.git synced 2024-11-09 15:02:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: directly resize images to 1024px

Browse source 
Former-commit-id: b6ed46be0e93bc735adc5a5c912893e5978ab2c5 [formerly 18f3b64ee956824024be1e4b13075ebb12b86176]
Former-commit-id: dbeaa87c8900e6191cf322bbac8994c6d9939aca
This commit is contained in:
Laurent Fainsin 2022-08-26 11:09:05 +02:00
parent 313d491143
commit 4696885a30
1 changed files with 11 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
src/data/dataset.py
View file

@ -1,6 +1,6 @@
import os
from pathlib import Path
import albumentations as A
import numpy as np
import torch
from PIL import Image
@ -16,6 +16,8 @@ class RealDataset(Dataset):
self.imgs = list(sorted(os.listdir(os.path.join(root, "images"))))
self.masks = list(sorted(os.listdir(os.path.join(root, "masks"))))
self.res = A.SmallestMaxSize(max_size=1024)
def __getitem__(self, idx):
# create paths from ids
image_path = os.path.join(self.root, "images", self.imgs[idx])
@ -23,11 +25,16 @@ class RealDataset(Dataset):
# load image and mask
image = Image.open(image_path).convert("RGB")
mask = Image.open(mask_path)
mask = Image.open(mask_path).convert("L")
# convert to numpy arrays
image = np.array(image)
mask = np.array(mask)
image = np.asarray(image)
mask = np.asarray(mask)
# resize images
aug = self.res(image=image, mask=mask)
image = aug["image"]
mask = aug["mask"]
# get ids from mask
obj_ids = np.unique(mask)