| import random |
| import numpy as np |
| from PIL import Image, ImageOps, ImageFilter |
| import torch |
| import torch.utils.data as data |
|
|
| __all__ = ['BaseDataset'] |
|
|
| class BaseDataset(data.Dataset): |
| def __init__(self, root, split, mode=None, transform=None, |
| target_transform=None, base_size=1024, crop_size=512): |
| self.root = root |
| self.transform = transform |
| self.target_transform = target_transform |
| self.split = split |
| self.mode = mode if mode is not None else split |
| self.base_size = base_size |
| self.crop_size = crop_size |
| if self.mode == 'train': |
| print('BaseDataset: base_size {}, crop_size {}'. \ |
| format(base_size, crop_size)) |
|
|
| @property |
| def num_class(self): |
| return self.NUM_CLASS |
|
|
| def _val_transform(self, img, mask): |
| outsize = self.crop_size |
| short_size = outsize |
| w, h = img.size |
| if w > h: |
| oh = short_size |
| ow = int(1.0 * w * oh / h) |
| else: |
| ow = short_size |
| oh = int(1.0 * h * ow / w) |
| img = img.resize((ow, oh), Image.BILINEAR) |
| mask = mask.resize((ow, oh), Image.NEAREST) |
| |
| w, h = img.size |
| x1 = int(round((w - outsize) / 2.)) |
| y1 = int(round((h - outsize) / 2.)) |
| img = img.crop((x1, y1, x1+outsize, y1+outsize)) |
| mask = mask.crop((x1, y1, x1+outsize, y1+outsize)) |
| |
| return img, self._mask_transform(mask) |
|
|
| def _testval_transform(self, img, mask): |
| outsize = self.crop_size |
| short_size = outsize |
| w, h = img.size |
| if w > h: |
| oh = short_size |
| ow = int(1.0 * w * oh / h) |
| else: |
| ow = short_size |
| oh = int(1.0 * h * ow / w) |
| img = img.resize((ow, oh), Image.BILINEAR) |
| return img, self._mask_transform(mask) |
|
|
| def _train_transform(self, img, mask): |
| |
| if random.random() < 0.5: |
| img = img.transpose(Image.FLIP_LEFT_RIGHT) |
| mask = mask.transpose(Image.FLIP_LEFT_RIGHT) |
| crop_size = self.crop_size |
| w, h = img.size |
| long_size = random.randint(int(self.base_size*0.5), int(self.base_size*2.0)) |
| if h > w: |
| oh = long_size |
| ow = int(1.0 * w * long_size / h + 0.5) |
| short_size = ow |
| else: |
| ow = long_size |
| oh = int(1.0 * h * long_size / w + 0.5) |
| short_size = oh |
| img = img.resize((ow, oh), Image.BILINEAR) |
| mask = mask.resize((ow, oh), Image.NEAREST) |
| |
| if short_size < crop_size: |
| padh = crop_size - oh if oh < crop_size else 0 |
| padw = crop_size - ow if ow < crop_size else 0 |
| img = ImageOps.expand(img, border=(0, 0, padw, padh), fill=0) |
| mask = ImageOps.expand(mask, border=(0, 0, padw, padh), fill=0) |
| |
| w, h = img.size |
| x1 = random.randint(0, w - crop_size) |
| y1 = random.randint(0, h - crop_size) |
| img = img.crop((x1, y1, x1+crop_size, y1+crop_size)) |
| mask = mask.crop((x1, y1, x1+crop_size, y1+crop_size)) |
| |
| return img, self._mask_transform(mask) |
|
|
| def _mask_transform(self, mask): |
| return torch.from_numpy(np.array(mask)).long() |
|
|
|
|
