feat: directly resize images to 1024px
Former-commit-id: b6ed46be0e93bc735adc5a5c912893e5978ab2c5 [formerly 18f3b64ee956824024be1e4b13075ebb12b86176] Former-commit-id: dbeaa87c8900e6191cf322bbac8994c6d9939aca
This commit is contained in:
parent
313d491143
commit
4696885a30
|
@ -1,6 +1,6 @@
|
||||||
import os
|
import os
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
|
import albumentations as A
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
@ -16,6 +16,8 @@ class RealDataset(Dataset):
|
||||||
self.imgs = list(sorted(os.listdir(os.path.join(root, "images"))))
|
self.imgs = list(sorted(os.listdir(os.path.join(root, "images"))))
|
||||||
self.masks = list(sorted(os.listdir(os.path.join(root, "masks"))))
|
self.masks = list(sorted(os.listdir(os.path.join(root, "masks"))))
|
||||||
|
|
||||||
|
self.res = A.SmallestMaxSize(max_size=1024)
|
||||||
|
|
||||||
def __getitem__(self, idx):
|
def __getitem__(self, idx):
|
||||||
# create paths from ids
|
# create paths from ids
|
||||||
image_path = os.path.join(self.root, "images", self.imgs[idx])
|
image_path = os.path.join(self.root, "images", self.imgs[idx])
|
||||||
|
@ -23,11 +25,16 @@ class RealDataset(Dataset):
|
||||||
|
|
||||||
# load image and mask
|
# load image and mask
|
||||||
image = Image.open(image_path).convert("RGB")
|
image = Image.open(image_path).convert("RGB")
|
||||||
mask = Image.open(mask_path)
|
mask = Image.open(mask_path).convert("L")
|
||||||
|
|
||||||
# convert to numpy arrays
|
# convert to numpy arrays
|
||||||
image = np.array(image)
|
image = np.asarray(image)
|
||||||
mask = np.array(mask)
|
mask = np.asarray(mask)
|
||||||
|
|
||||||
|
# resize images
|
||||||
|
aug = self.res(image=image, mask=mask)
|
||||||
|
image = aug["image"]
|
||||||
|
mask = aug["mask"]
|
||||||
|
|
||||||
# get ids from mask
|
# get ids from mask
|
||||||
obj_ids = np.unique(mask)
|
obj_ids = np.unique(mask)
|
||||||
|
|
Loading…
Reference in a new issue