| import os |
| import sys |
| import random |
| import numpy as np |
| from tqdm import tqdm, trange |
| from PIL import Image, ImageOps, ImageFilter |
|
|
| import torch |
| import torch.utils.data as data |
| import torchvision.transforms as transform |
|
|
| from datasets.base import BaseDataset |
|
|
| class CitySegmentation(BaseDataset): |
| NUM_CLASS = 19 |
| def __init__(self, root, split='val', mode='testval', transform=None, target_transform=None, **kwargs): |
| super(CitySegmentation, self).__init__( |
| root, split, mode, transform, target_transform, **kwargs) |
| self.images, self.mask_paths = get_city_pairs(self.root, self.split) |
| assert (len(self.images) == len(self.mask_paths)) |
| if len(self.images) == 0: |
| raise RuntimeError("Found 0 images in subfolders of: \ |
| " + self.root + "\n") |
| self._indices = np.array(range(-1, 19)) |
| self._classes = np.array([0, 7, 8, 11, 12, 13, 17, 19, 20, 21, 22, |
| 23, 24, 25, 26, 27, 28, 31, 32, 33]) |
| self._key = np.array([-1, -1, -1, -1, -1, -1, |
| -1, -1, 0, 1, -1, -1, |
| 2, 3, 4, -1, -1, -1, |
| 5, -1, 6, 7, 8, 9, |
| 10, 11, 12, 13, 14, 15, |
| -1, -1, 16, 17, 18]) |
| self._mapping = np.array(range(-1, len(self._key)-1)).astype('int32') |
|
|
| def _class_to_index(self, mask): |
| |
| values = np.unique(mask) |
| for i in range(len(values)): |
| assert(values[i] in self._mapping) |
| index = np.digitize(mask.ravel(), self._mapping, right=True) |
| return self._key[index].reshape(mask.shape) |
|
|
| def __getitem__(self, index): |
| img = Image.open(self.images[index]).convert('RGB') |
| mask = Image.open(self.mask_paths[index]) |
| if self.mode == 'testval': |
| img, mask = self._testval_transform(img, mask) |
| elif self.mode == 'val': |
| img, mask = self._val_transform(img, mask) |
| elif self.mode == 'train': |
| img, mask = self._train_transform(img, mask) |
|
|
| if self.transform is not None: |
| img = self.transform(img) |
| if self.target_transform is not None: |
| mask = self.target_transform(mask) |
| return img, mask |
|
|
| def _mask_transform(self, mask): |
| target = self._class_to_index(np.array(mask).astype('int32')) |
| return torch.from_numpy(target).long() |
|
|
| def __len__(self): |
| return len(self.images) |
|
|
|
|
| def get_city_pairs(folder, split='val'): |
| def get_path_pairs(img_folder, mask_folder): |
| img_paths = [] |
| mask_paths = [] |
| for root, directories, files in os.walk(img_folder): |
| for filename in files: |
| if filename.endswith(".png"): |
| imgpath = os.path.join(root, filename) |
| foldername = os.path.basename(os.path.dirname(imgpath)) |
| maskname = filename.replace('leftImg8bit','gtFine_labelIds') |
| maskpath = os.path.join(mask_folder, foldername, maskname) |
| if os.path.isfile(imgpath) and os.path.isfile(maskpath): |
| img_paths.append(imgpath) |
| mask_paths.append(maskpath) |
| else: |
| print('cannot find the mask or image:', imgpath, maskpath) |
| print('Found {} images in the folder {}'.format(len(img_paths), img_folder)) |
| return img_paths, mask_paths |
|
|
| img_folder = os.path.join(folder, 'leftImg8bit/' + split) |
| mask_folder = os.path.join(folder, 'gtFine/'+ split) |
| img_paths, mask_paths = get_path_pairs(img_folder, mask_folder) |
| return img_paths, mask_paths |
|
|
