358 lines
13 KiB
Python
358 lines
13 KiB
Python
import random
|
|
import cv2
|
|
import numpy as np
|
|
import numbers
|
|
import collections
|
|
|
|
# copy from: https://github.com/cardwing/Codes-for-Lane-Detection/blob/master/ERFNet-CULane-PyTorch/utils/transforms.py
|
|
|
|
__all__ = ['GroupRandomCrop', 'GroupCenterCrop', 'GroupRandomPad', 'GroupCenterPad',
|
|
'GroupRandomScale', 'GroupRandomHorizontalFlip', 'GroupNormalize']
|
|
|
|
|
|
class SampleResize(object):
|
|
def __init__(self, size):
|
|
assert (isinstance(size, collections.Iterable) and len(size) == 2)
|
|
self.size = size
|
|
|
|
def __call__(self, sample):
|
|
out = list()
|
|
out.append(cv2.resize(sample[0], self.size,
|
|
interpolation=cv2.INTER_CUBIC))
|
|
if len(sample) > 1:
|
|
out.append(cv2.resize(sample[1], self.size,
|
|
interpolation=cv2.INTER_NEAREST))
|
|
return out
|
|
|
|
|
|
class GroupRandomCrop(object):
|
|
def __init__(self, size):
|
|
if isinstance(size, numbers.Number):
|
|
self.size = (int(size), int(size))
|
|
else:
|
|
self.size = size
|
|
|
|
def __call__(self, img_group):
|
|
h, w = img_group[0].shape[0:2]
|
|
th, tw = self.size
|
|
|
|
out_images = list()
|
|
h1 = random.randint(0, max(0, h - th))
|
|
w1 = random.randint(0, max(0, w - tw))
|
|
h2 = min(h1 + th, h)
|
|
w2 = min(w1 + tw, w)
|
|
|
|
for img in img_group:
|
|
assert (img.shape[0] == h and img.shape[1] == w)
|
|
out_images.append(img[h1:h2, w1:w2, ...])
|
|
return out_images
|
|
|
|
|
|
class GroupRandomCropRatio(object):
|
|
def __init__(self, size):
|
|
if isinstance(size, numbers.Number):
|
|
self.size = (int(size), int(size))
|
|
else:
|
|
self.size = size
|
|
|
|
def __call__(self, img_group):
|
|
h, w = img_group[0].shape[0:2]
|
|
tw, th = self.size
|
|
|
|
out_images = list()
|
|
h1 = random.randint(0, max(0, h - th))
|
|
w1 = random.randint(0, max(0, w - tw))
|
|
h2 = min(h1 + th, h)
|
|
w2 = min(w1 + tw, w)
|
|
|
|
for img in img_group:
|
|
assert (img.shape[0] == h and img.shape[1] == w)
|
|
out_images.append(img[h1:h2, w1:w2, ...])
|
|
return out_images
|
|
|
|
|
|
class GroupCenterCrop(object):
|
|
def __init__(self, size):
|
|
if isinstance(size, numbers.Number):
|
|
self.size = (int(size), int(size))
|
|
else:
|
|
self.size = size
|
|
|
|
def __call__(self, img_group):
|
|
h, w = img_group[0].shape[0:2]
|
|
th, tw = self.size
|
|
|
|
out_images = list()
|
|
h1 = max(0, int((h - th) / 2))
|
|
w1 = max(0, int((w - tw) / 2))
|
|
h2 = min(h1 + th, h)
|
|
w2 = min(w1 + tw, w)
|
|
|
|
for img in img_group:
|
|
assert (img.shape[0] == h and img.shape[1] == w)
|
|
out_images.append(img[h1:h2, w1:w2, ...])
|
|
return out_images
|
|
|
|
|
|
class GroupRandomPad(object):
|
|
def __init__(self, size, padding):
|
|
if isinstance(size, numbers.Number):
|
|
self.size = (int(size), int(size))
|
|
else:
|
|
self.size = size
|
|
self.padding = padding
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.padding) == len(img_group))
|
|
h, w = img_group[0].shape[0:2]
|
|
th, tw = self.size
|
|
|
|
out_images = list()
|
|
h1 = random.randint(0, max(0, th - h))
|
|
w1 = random.randint(0, max(0, tw - w))
|
|
h2 = max(th - h - h1, 0)
|
|
w2 = max(tw - w - w1, 0)
|
|
|
|
for img, padding in zip(img_group, self.padding):
|
|
assert (img.shape[0] == h and img.shape[1] == w)
|
|
out_images.append(cv2.copyMakeBorder(
|
|
img, h1, h2, w1, w2, cv2.BORDER_CONSTANT, value=padding))
|
|
if len(img.shape) > len(out_images[-1].shape):
|
|
out_images[-1] = out_images[-1][...,
|
|
np.newaxis] # single channel image
|
|
return out_images
|
|
|
|
|
|
class GroupCenterPad(object):
|
|
def __init__(self, size, padding):
|
|
if isinstance(size, numbers.Number):
|
|
self.size = (int(size), int(size))
|
|
else:
|
|
self.size = size
|
|
self.padding = padding
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.padding) == len(img_group))
|
|
h, w = img_group[0].shape[0:2]
|
|
th, tw = self.size
|
|
|
|
out_images = list()
|
|
h1 = max(0, int((th - h) / 2))
|
|
w1 = max(0, int((tw - w) / 2))
|
|
h2 = max(th - h - h1, 0)
|
|
w2 = max(tw - w - w1, 0)
|
|
|
|
for img, padding in zip(img_group, self.padding):
|
|
assert (img.shape[0] == h and img.shape[1] == w)
|
|
out_images.append(cv2.copyMakeBorder(
|
|
img, h1, h2, w1, w2, cv2.BORDER_CONSTANT, value=padding))
|
|
if len(img.shape) > len(out_images[-1].shape):
|
|
out_images[-1] = out_images[-1][...,
|
|
np.newaxis] # single channel image
|
|
return out_images
|
|
|
|
|
|
class GroupConcerPad(object):
|
|
def __init__(self, size, padding):
|
|
if isinstance(size, numbers.Number):
|
|
self.size = (int(size), int(size))
|
|
else:
|
|
self.size = size
|
|
self.padding = padding
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.padding) == len(img_group))
|
|
h, w = img_group[0].shape[0:2]
|
|
th, tw = self.size
|
|
|
|
out_images = list()
|
|
h1 = 0
|
|
w1 = 0
|
|
h2 = max(th - h - h1, 0)
|
|
w2 = max(tw - w - w1, 0)
|
|
|
|
for img, padding in zip(img_group, self.padding):
|
|
assert (img.shape[0] == h and img.shape[1] == w)
|
|
out_images.append(cv2.copyMakeBorder(
|
|
img, h1, h2, w1, w2, cv2.BORDER_CONSTANT, value=padding))
|
|
if len(img.shape) > len(out_images[-1].shape):
|
|
out_images[-1] = out_images[-1][...,
|
|
np.newaxis] # single channel image
|
|
return out_images
|
|
|
|
|
|
class GroupRandomScaleNew(object):
|
|
def __init__(self, size=(976, 208), interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)):
|
|
self.size = size
|
|
self.interpolation = interpolation
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.interpolation) == len(img_group))
|
|
scale_w, scale_h = self.size[0] * 1.0 / 1640, self.size[1] * 1.0 / 590
|
|
out_images = list()
|
|
for img, interpolation in zip(img_group, self.interpolation):
|
|
out_images.append(cv2.resize(img, None, fx=scale_w,
|
|
fy=scale_h, interpolation=interpolation))
|
|
if len(img.shape) > len(out_images[-1].shape):
|
|
out_images[-1] = out_images[-1][...,
|
|
np.newaxis] # single channel image
|
|
return out_images
|
|
|
|
|
|
class GroupRandomScale(object):
|
|
def __init__(self, size=(0.5, 1.5), interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)):
|
|
self.size = size
|
|
self.interpolation = interpolation
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.interpolation) == len(img_group))
|
|
scale = random.uniform(self.size[0], self.size[1])
|
|
out_images = list()
|
|
for img, interpolation in zip(img_group, self.interpolation):
|
|
out_images.append(cv2.resize(img, None, fx=scale,
|
|
fy=scale, interpolation=interpolation))
|
|
if len(img.shape) > len(out_images[-1].shape):
|
|
out_images[-1] = out_images[-1][...,
|
|
np.newaxis] # single channel image
|
|
return out_images
|
|
|
|
|
|
class GroupRandomMultiScale(object):
|
|
def __init__(self, size=(0.5, 1.5), interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)):
|
|
self.size = size
|
|
self.interpolation = interpolation
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.interpolation) == len(img_group))
|
|
scales = [0.5, 1.0, 1.5] # random.uniform(self.size[0], self.size[1])
|
|
out_images = list()
|
|
for scale in scales:
|
|
for img, interpolation in zip(img_group, self.interpolation):
|
|
out_images.append(cv2.resize(
|
|
img, None, fx=scale, fy=scale, interpolation=interpolation))
|
|
if len(img.shape) > len(out_images[-1].shape):
|
|
out_images[-1] = out_images[-1][...,
|
|
np.newaxis] # single channel image
|
|
return out_images
|
|
|
|
|
|
class GroupRandomScaleRatio(object):
|
|
def __init__(self, size=(680, 762, 562, 592), interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)):
|
|
self.size = size
|
|
self.interpolation = interpolation
|
|
self.origin_id = [0, 1360, 580, 768, 255, 300, 680, 710, 312, 1509, 800, 1377, 880, 910, 1188, 128, 960, 1784,
|
|
1414, 1150, 512, 1162, 950, 750, 1575, 708, 2111, 1848, 1071, 1204, 892, 639, 2040, 1524, 832, 1122, 1224, 2295]
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.interpolation) == len(img_group))
|
|
w_scale = random.randint(self.size[0], self.size[1])
|
|
h_scale = random.randint(self.size[2], self.size[3])
|
|
h, w, _ = img_group[0].shape
|
|
out_images = list()
|
|
out_images.append(cv2.resize(img_group[0], None, fx=w_scale*1.0/w, fy=h_scale*1.0/h,
|
|
interpolation=self.interpolation[0])) # fx=w_scale*1.0/w, fy=h_scale*1.0/h
|
|
### process label map ###
|
|
origin_label = cv2.resize(
|
|
img_group[1], None, fx=w_scale*1.0/w, fy=h_scale*1.0/h, interpolation=self.interpolation[1])
|
|
origin_label = origin_label.astype(int)
|
|
label = origin_label[:, :, 0] * 5 + \
|
|
origin_label[:, :, 1] * 3 + origin_label[:, :, 2]
|
|
new_label = np.ones(label.shape) * 100
|
|
new_label = new_label.astype(int)
|
|
for cnt in range(37):
|
|
new_label = (
|
|
label == self.origin_id[cnt]) * (cnt - 100) + new_label
|
|
new_label = (label == self.origin_id[37]) * (36 - 100) + new_label
|
|
assert(100 not in np.unique(new_label))
|
|
out_images.append(new_label)
|
|
return out_images
|
|
|
|
|
|
class GroupRandomRotation(object):
|
|
def __init__(self, degree=(-10, 10), interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST), padding=None):
|
|
self.degree = degree
|
|
self.interpolation = interpolation
|
|
self.padding = padding
|
|
if self.padding is None:
|
|
self.padding = [0, 0]
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.interpolation) == len(img_group))
|
|
v = random.random()
|
|
if v < 0.5:
|
|
degree = random.uniform(self.degree[0], self.degree[1])
|
|
h, w = img_group[0].shape[0:2]
|
|
center = (w / 2, h / 2)
|
|
map_matrix = cv2.getRotationMatrix2D(center, degree, 1.0)
|
|
out_images = list()
|
|
for img, interpolation, padding in zip(img_group, self.interpolation, self.padding):
|
|
out_images.append(cv2.warpAffine(
|
|
img, map_matrix, (w, h), flags=interpolation, borderMode=cv2.BORDER_CONSTANT, borderValue=padding))
|
|
if len(img.shape) > len(out_images[-1].shape):
|
|
out_images[-1] = out_images[-1][...,
|
|
np.newaxis] # single channel image
|
|
return out_images
|
|
else:
|
|
return img_group
|
|
|
|
|
|
class GroupRandomBlur(object):
|
|
def __init__(self, applied):
|
|
self.applied = applied
|
|
|
|
def __call__(self, img_group):
|
|
assert (len(self.applied) == len(img_group))
|
|
v = random.random()
|
|
if v < 0.5:
|
|
out_images = []
|
|
for img, a in zip(img_group, self.applied):
|
|
if a:
|
|
img = cv2.GaussianBlur(
|
|
img, (5, 5), random.uniform(1e-6, 0.6))
|
|
out_images.append(img)
|
|
if len(img.shape) > len(out_images[-1].shape):
|
|
out_images[-1] = out_images[-1][...,
|
|
np.newaxis] # single channel image
|
|
return out_images
|
|
else:
|
|
return img_group
|
|
|
|
|
|
class GroupRandomHorizontalFlip(object):
|
|
"""Randomly horizontally flips the given numpy Image with a probability of 0.5
|
|
"""
|
|
|
|
def __init__(self, is_flow=False):
|
|
self.is_flow = is_flow
|
|
|
|
def __call__(self, img_group, is_flow=False):
|
|
v = random.random()
|
|
if v < 0.5:
|
|
out_images = [np.fliplr(img) for img in img_group]
|
|
if self.is_flow:
|
|
for i in range(0, len(out_images), 2):
|
|
# invert flow pixel values when flipping
|
|
out_images[i] = -out_images[i]
|
|
return out_images
|
|
else:
|
|
return img_group
|
|
|
|
|
|
class GroupNormalize(object):
|
|
def __init__(self, mean, std):
|
|
self.mean = mean
|
|
self.std = std
|
|
|
|
def __call__(self, img_group):
|
|
out_images = list()
|
|
for img, m, s in zip(img_group, self.mean, self.std):
|
|
if len(m) == 1:
|
|
img = img - np.array(m) # single channel image
|
|
img = img / np.array(s)
|
|
else:
|
|
img = img - np.array(m)[np.newaxis, np.newaxis, ...]
|
|
img = img / np.array(s)[np.newaxis, np.newaxis, ...]
|
|
out_images.append(img)
|
|
|
|
return out_images
|