| | 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 |
| |
|
