You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
246 lines
8.7 KiB
246 lines
8.7 KiB
import math
|
|
import numbers
|
|
import random
|
|
|
|
from PIL import Image, ImageOps
|
|
import numpy as np
|
|
|
|
|
|
class Compose(object):
|
|
def __init__(self, transforms):
|
|
self.transforms = transforms
|
|
|
|
def __call__(self, img, mask):
|
|
assert img.size == mask.size
|
|
for t in self.transforms:
|
|
img, mask = t(img, mask)
|
|
return img, mask
|
|
|
|
|
|
class RandomCrop(object):
|
|
def __init__(self, size, padding=0):
|
|
if isinstance(size, numbers.Number):
|
|
self.size = (int(size), int(size))
|
|
else:
|
|
self.size = size
|
|
self.padding = padding
|
|
|
|
def __call__(self, img, mask):
|
|
if self.padding > 0:
|
|
img = ImageOps.expand(img, border=self.padding, fill=0)
|
|
mask = ImageOps.expand(mask, border=self.padding, fill=0)
|
|
|
|
assert img.size == mask.size
|
|
w, h = img.size
|
|
th, tw = self.size
|
|
if w == tw and h == th:
|
|
return img, mask
|
|
if w < tw or h < th:
|
|
return img.resize((tw, th), Image.BILINEAR), mask.resize((tw, th), Image.NEAREST)
|
|
|
|
x1 = random.randint(0, w - tw)
|
|
y1 = random.randint(0, h - th)
|
|
return img.crop((x1, y1, x1 + tw, y1 + th)), mask.crop((x1, y1, x1 + tw, y1 + th))
|
|
|
|
|
|
class CenterCrop(object):
|
|
def __init__(self, size):
|
|
if isinstance(size, numbers.Number):
|
|
self.size = (int(size), int(size))
|
|
else:
|
|
self.size = size
|
|
|
|
def __call__(self, img, mask):
|
|
assert img.size == mask.size
|
|
w, h = img.size
|
|
th, tw = self.size
|
|
x1 = int(round((w - tw) / 2.))
|
|
y1 = int(round((h - th) / 2.))
|
|
return img.crop((x1, y1, x1 + tw, y1 + th)), mask.crop((x1, y1, x1 + tw, y1 + th))
|
|
|
|
|
|
class RandomHorizontallyFlip(object):
|
|
def __call__(self, img, mask):
|
|
if random.random() < 0.5:
|
|
return img.transpose(Image.FLIP_LEFT_RIGHT), mask.transpose(Image.FLIP_LEFT_RIGHT)
|
|
return img, mask
|
|
|
|
|
|
class FreeScale(object):
|
|
def __init__(self, size):
|
|
self.size = tuple(reversed(size)) # size: (h, w)
|
|
|
|
def __call__(self, img, mask):
|
|
assert img.size == mask.size
|
|
return img.resize(self.size, Image.BILINEAR), mask.resize(self.size, Image.NEAREST)
|
|
|
|
|
|
class Scale(object):
|
|
def __init__(self, size):
|
|
self.size = size
|
|
|
|
def __call__(self, img, mask):
|
|
assert img.size == mask.size
|
|
w, h = img.size
|
|
if (w >= h and w == self.size) or (h >= w and h == self.size):
|
|
return img, mask
|
|
if w > h:
|
|
ow = self.size
|
|
oh = int(self.size * h / w)
|
|
return img.resize((ow, oh), Image.BILINEAR), mask.resize((ow, oh), Image.NEAREST)
|
|
else:
|
|
oh = self.size
|
|
ow = int(self.size * w / h)
|
|
return img.resize((ow, oh), Image.BILINEAR), mask.resize((ow, oh), Image.NEAREST)
|
|
|
|
|
|
class RandomSizedCrop(object):
|
|
def __init__(self, size):
|
|
self.size = size
|
|
|
|
def __call__(self, img, mask):
|
|
assert img.size == mask.size
|
|
for attempt in range(10):
|
|
area = img.size[0] * img.size[1]
|
|
target_area = random.uniform(0.45, 1.0) * area
|
|
aspect_ratio = random.uniform(0.5, 2)
|
|
|
|
w = int(round(math.sqrt(target_area * aspect_ratio)))
|
|
h = int(round(math.sqrt(target_area / aspect_ratio)))
|
|
|
|
if random.random() < 0.5:
|
|
w, h = h, w
|
|
|
|
if w <= img.size[0] and h <= img.size[1]:
|
|
x1 = random.randint(0, img.size[0] - w)
|
|
y1 = random.randint(0, img.size[1] - h)
|
|
|
|
img = img.crop((x1, y1, x1 + w, y1 + h))
|
|
mask = mask.crop((x1, y1, x1 + w, y1 + h))
|
|
assert (img.size == (w, h))
|
|
|
|
return img.resize((self.size, self.size), Image.BILINEAR), mask.resize((self.size, self.size),
|
|
Image.NEAREST)
|
|
|
|
# Fallback
|
|
scale = Scale(self.size)
|
|
crop = CenterCrop(self.size)
|
|
return crop(*scale(img, mask))
|
|
|
|
|
|
class RandomRotate(object):
|
|
def __init__(self, degree):
|
|
self.degree = degree
|
|
|
|
def __call__(self, img, mask):
|
|
rotate_degree = random.random() * 2 * self.degree - self.degree
|
|
return img.rotate(rotate_degree, Image.BILINEAR), mask.rotate(rotate_degree, Image.NEAREST)
|
|
|
|
|
|
class RandomSized(object):
|
|
def __init__(self, size):
|
|
self.size = size
|
|
self.scale = Scale(self.size)
|
|
self.crop = RandomCrop(self.size)
|
|
|
|
def __call__(self, img, mask):
|
|
assert img.size == mask.size
|
|
|
|
w = int(random.uniform(0.5, 2) * img.size[0])
|
|
h = int(random.uniform(0.5, 2) * img.size[1])
|
|
|
|
img, mask = img.resize((w, h), Image.BILINEAR), mask.resize((w, h), Image.NEAREST)
|
|
|
|
return self.crop(*self.scale(img, mask))
|
|
|
|
|
|
class SlidingCropOld(object):
|
|
def __init__(self, crop_size, stride_rate, ignore_label):
|
|
self.crop_size = crop_size
|
|
self.stride_rate = stride_rate
|
|
self.ignore_label = ignore_label
|
|
|
|
def _pad(self, img, mask):
|
|
h, w = img.shape[: 2]
|
|
pad_h = max(self.crop_size - h, 0)
|
|
pad_w = max(self.crop_size - w, 0)
|
|
img = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), 'constant')
|
|
mask = np.pad(mask, ((0, pad_h), (0, pad_w)), 'constant', constant_values=self.ignore_label)
|
|
return img, mask
|
|
|
|
def __call__(self, img, mask):
|
|
assert img.size == mask.size
|
|
|
|
w, h = img.size
|
|
long_size = max(h, w)
|
|
|
|
img = np.array(img)
|
|
mask = np.array(mask)
|
|
|
|
if long_size > self.crop_size:
|
|
stride = int(math.ceil(self.crop_size * self.stride_rate))
|
|
h_step_num = int(math.ceil((h - self.crop_size) / float(stride))) + 1
|
|
w_step_num = int(math.ceil((w - self.crop_size) / float(stride))) + 1
|
|
img_sublist, mask_sublist = [], []
|
|
for yy in xrange(h_step_num):
|
|
for xx in xrange(w_step_num):
|
|
sy, sx = yy * stride, xx * stride
|
|
ey, ex = sy + self.crop_size, sx + self.crop_size
|
|
img_sub = img[sy: ey, sx: ex, :]
|
|
mask_sub = mask[sy: ey, sx: ex]
|
|
img_sub, mask_sub = self._pad(img_sub, mask_sub)
|
|
img_sublist.append(Image.fromarray(img_sub.astype(np.uint8)).convert('RGB'))
|
|
mask_sublist.append(Image.fromarray(mask_sub.astype(np.uint8)).convert('P'))
|
|
return img_sublist, mask_sublist
|
|
else:
|
|
img, mask = self._pad(img, mask)
|
|
img = Image.fromarray(img.astype(np.uint8)).convert('RGB')
|
|
mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
|
|
return img, mask
|
|
|
|
|
|
class SlidingCrop(object):
|
|
def __init__(self, crop_size, stride_rate, ignore_label):
|
|
self.crop_size = crop_size
|
|
self.stride_rate = stride_rate
|
|
self.ignore_label = ignore_label
|
|
|
|
def _pad(self, img, mask):
|
|
h, w = img.shape[: 2]
|
|
pad_h = max(self.crop_size - h, 0)
|
|
pad_w = max(self.crop_size - w, 0)
|
|
img = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), 'constant')
|
|
mask = np.pad(mask, ((0, pad_h), (0, pad_w)), 'constant', constant_values=self.ignore_label)
|
|
return img, mask, h, w
|
|
|
|
def __call__(self, img, mask):
|
|
assert img.size == mask.size
|
|
|
|
w, h = img.size
|
|
long_size = max(h, w)
|
|
|
|
img = np.array(img)
|
|
mask = np.array(mask)
|
|
|
|
if long_size > self.crop_size:
|
|
stride = int(math.ceil(self.crop_size * self.stride_rate))
|
|
h_step_num = int(math.ceil((h - self.crop_size) / float(stride))) + 1
|
|
w_step_num = int(math.ceil((w - self.crop_size) / float(stride))) + 1
|
|
img_slices, mask_slices, slices_info = [], [], []
|
|
for yy in range(h_step_num):
|
|
for xx in range(w_step_num):
|
|
sy, sx = yy * stride, xx * stride
|
|
ey, ex = sy + self.crop_size, sx + self.crop_size
|
|
img_sub = img[sy: ey, sx: ex, :]
|
|
mask_sub = mask[sy: ey, sx: ex]
|
|
img_sub, mask_sub, sub_h, sub_w = self._pad(img_sub, mask_sub)
|
|
img_slices.append(Image.fromarray(img_sub.astype(np.uint8)).convert('RGB'))
|
|
mask_slices.append(Image.fromarray(mask_sub.astype(np.uint8)).convert('P'))
|
|
slices_info.append([sy, ey, sx, ex, sub_h, sub_w])
|
|
return img_slices, mask_slices, slices_info
|
|
else:
|
|
img, mask, sub_h, sub_w = self._pad(img, mask)
|
|
img = Image.fromarray(img.astype(np.uint8)).convert('RGB')
|
|
mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
|
|
return [img], [mask], [[0, sub_h, 0, sub_w, sub_h, sub_w]]
|