Spaces:
Build error
Build error
File size: 2,291 Bytes
6f774ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import os
from torch.utils.data import Dataset
from PIL import Image
import torchvision.transforms as transforms
import numpy as np
class AerialImageDataset(Dataset):
def __init__(self, image_dir, mask_dir, transform=None):
self.image_dir = image_dir
self.mask_dir = mask_dir
self.transform = transform
self.images = os.listdir(self.image_dir)
self.Hex_Classes = [
('Unlabeled', '#FF2285'),
('Early Blight','#00FCC7'),
('Late Blight', '#5600FE'),
('Leaf Minor', '#000000')
]
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
img_path = os.path.join(self.image_dir, self.images[idx])
mask_path = os.path.join(self.mask_dir, self.images[idx].replace('.jpg', '.png'))
image = Image.open(img_path)
mask = Image.open(mask_path).resize((512, 512))
# print(mask.size)
# print(mask.size)
mask = np.array(mask)
mask = self.encode_segmap(mask)
mask = mask.astype(np.uint8) # Convert data type to uint8
# print(mask.shape)
mask = Image.fromarray(mask) # Convert mask -> PIL
if self.transform:
image = self.transform(image)
mask = self.transform(mask)
return image, mask
def encode_segmap(self, mask):
mask = mask.astype(int)
label_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int16) # height, width -> 0
for i, (name, color) in enumerate(self.Hex_Classes):
if mask.ndim == 3:
label_mask[(mask[:,:,0] == int(color[1:3], 16)) & (mask[:,:,1] == int(color[3:5], 16)) & (mask[:,:,2] == int(color[5:7], 16))] = i
elif mask.ndim == 2:
label_mask[(mask == int(color[1:3], 16))] = i
# print("Warning ndim = 2")
# return None
msk = np.zeros((512,512,4))
for i in [0,1,2,3]:
if i == 0:
msk_ind = np.where(label_mask == i, 4, 0)
msk[:,:,i] = msk_ind
else:
msk_ind = np.where(label_mask == i, i, 0)
msk[:,:,i] = msk_ind
# print("mask shape",type(msk))
return msk |