Spaces:
Running
on
Zero
Running
on
Zero
""" | |
Copyright (c) 2022, salesforce.com, inc. | |
All rights reserved. | |
SPDX-License-Identifier: BSD-3-Clause | |
For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause | |
""" | |
import cv2 | |
import numpy as np | |
import torch | |
## aug functions | |
def identity_func(img): | |
return img | |
def autocontrast_func(img, cutoff=0): | |
""" | |
same output as PIL.ImageOps.autocontrast | |
""" | |
n_bins = 256 | |
def tune_channel(ch): | |
n = ch.size | |
cut = cutoff * n // 100 | |
if cut == 0: | |
high, low = ch.max(), ch.min() | |
else: | |
hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins]) | |
low = np.argwhere(np.cumsum(hist) > cut) | |
low = 0 if low.shape[0] == 0 else low[0] | |
high = np.argwhere(np.cumsum(hist[::-1]) > cut) | |
high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0] | |
if high <= low: | |
table = np.arange(n_bins) | |
else: | |
scale = (n_bins - 1) / (high - low) | |
offset = -low * scale | |
table = np.arange(n_bins) * scale + offset | |
table[table < 0] = 0 | |
table[table > n_bins - 1] = n_bins - 1 | |
table = table.clip(0, 255).astype(np.uint8) | |
return table[ch] | |
channels = [tune_channel(ch) for ch in cv2.split(img)] | |
out = cv2.merge(channels) | |
return out | |
def equalize_func(img): | |
""" | |
same output as PIL.ImageOps.equalize | |
PIL's implementation is different from cv2.equalize | |
""" | |
n_bins = 256 | |
def tune_channel(ch): | |
hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins]) | |
non_zero_hist = hist[hist != 0].reshape(-1) | |
step = np.sum(non_zero_hist[:-1]) // (n_bins - 1) | |
if step == 0: | |
return ch | |
n = np.empty_like(hist) | |
n[0] = step // 2 | |
n[1:] = hist[:-1] | |
table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8) | |
return table[ch] | |
channels = [tune_channel(ch) for ch in cv2.split(img)] | |
out = cv2.merge(channels) | |
return out | |
def rotate_func(img, degree, fill=(0, 0, 0)): | |
""" | |
like PIL, rotate by degree, not radians | |
""" | |
H, W = img.shape[0], img.shape[1] | |
center = W / 2, H / 2 | |
M = cv2.getRotationMatrix2D(center, degree, 1) | |
out = cv2.warpAffine(img, M, (W, H), borderValue=fill) | |
return out | |
def solarize_func(img, thresh=128): | |
""" | |
same output as PIL.ImageOps.posterize | |
""" | |
table = np.array([el if el < thresh else 255 - el for el in range(256)]) | |
table = table.clip(0, 255).astype(np.uint8) | |
out = table[img] | |
return out | |
def color_func(img, factor): | |
""" | |
same output as PIL.ImageEnhance.Color | |
""" | |
## implementation according to PIL definition, quite slow | |
# degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis] | |
# out = blend(degenerate, img, factor) | |
# M = ( | |
# np.eye(3) * factor | |
# + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor) | |
# )[np.newaxis, np.newaxis, :] | |
M = np.float32( | |
[[0.886, -0.114, -0.114], [-0.587, 0.413, -0.587], [-0.299, -0.299, 0.701]] | |
) * factor + np.float32([[0.114], [0.587], [0.299]]) | |
out = np.matmul(img, M).clip(0, 255).astype(np.uint8) | |
return out | |
def contrast_func(img, factor): | |
""" | |
same output as PIL.ImageEnhance.Contrast | |
""" | |
mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299])) | |
table = ( | |
np.array([(el - mean) * factor + mean for el in range(256)]) | |
.clip(0, 255) | |
.astype(np.uint8) | |
) | |
out = table[img] | |
return out | |
def brightness_func(img, factor): | |
""" | |
same output as PIL.ImageEnhance.Contrast | |
""" | |
table = (np.arange(256, dtype=np.float32) * factor).clip(0, 255).astype(np.uint8) | |
out = table[img] | |
return out | |
def sharpness_func(img, factor): | |
""" | |
The differences the this result and PIL are all on the 4 boundaries, the center | |
areas are same | |
""" | |
kernel = np.ones((3, 3), dtype=np.float32) | |
kernel[1][1] = 5 | |
kernel /= 13 | |
degenerate = cv2.filter2D(img, -1, kernel) | |
if factor == 0.0: | |
out = degenerate | |
elif factor == 1.0: | |
out = img | |
else: | |
out = img.astype(np.float32) | |
degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :] | |
out[1:-1, 1:-1, :] = degenerate + factor * (out[1:-1, 1:-1, :] - degenerate) | |
out = out.astype(np.uint8) | |
return out | |
def shear_x_func(img, factor, fill=(0, 0, 0)): | |
H, W = img.shape[0], img.shape[1] | |
M = np.float32([[1, factor, 0], [0, 1, 0]]) | |
out = cv2.warpAffine( | |
img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR | |
).astype(np.uint8) | |
return out | |
def translate_x_func(img, offset, fill=(0, 0, 0)): | |
""" | |
same output as PIL.Image.transform | |
""" | |
H, W = img.shape[0], img.shape[1] | |
M = np.float32([[1, 0, -offset], [0, 1, 0]]) | |
out = cv2.warpAffine( | |
img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR | |
).astype(np.uint8) | |
return out | |
def translate_y_func(img, offset, fill=(0, 0, 0)): | |
""" | |
same output as PIL.Image.transform | |
""" | |
H, W = img.shape[0], img.shape[1] | |
M = np.float32([[1, 0, 0], [0, 1, -offset]]) | |
out = cv2.warpAffine( | |
img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR | |
).astype(np.uint8) | |
return out | |
def posterize_func(img, bits): | |
""" | |
same output as PIL.ImageOps.posterize | |
""" | |
out = np.bitwise_and(img, np.uint8(255 << (8 - bits))) | |
return out | |
def shear_y_func(img, factor, fill=(0, 0, 0)): | |
H, W = img.shape[0], img.shape[1] | |
M = np.float32([[1, 0, 0], [factor, 1, 0]]) | |
out = cv2.warpAffine( | |
img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR | |
).astype(np.uint8) | |
return out | |
def cutout_func(img, pad_size, replace=(0, 0, 0)): | |
replace = np.array(replace, dtype=np.uint8) | |
H, W = img.shape[0], img.shape[1] | |
rh, rw = np.random.random(2) | |
pad_size = pad_size // 2 | |
ch, cw = int(rh * H), int(rw * W) | |
x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H) | |
y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W) | |
out = img.copy() | |
out[x1:x2, y1:y2, :] = replace | |
return out | |
### level to args | |
def enhance_level_to_args(MAX_LEVEL): | |
def level_to_args(level): | |
return ((level / MAX_LEVEL) * 1.8 + 0.1,) | |
return level_to_args | |
def shear_level_to_args(MAX_LEVEL, replace_value): | |
def level_to_args(level): | |
level = (level / MAX_LEVEL) * 0.3 | |
if np.random.random() > 0.5: | |
level = -level | |
return (level, replace_value) | |
return level_to_args | |
def translate_level_to_args(translate_const, MAX_LEVEL, replace_value): | |
def level_to_args(level): | |
level = (level / MAX_LEVEL) * float(translate_const) | |
if np.random.random() > 0.5: | |
level = -level | |
return (level, replace_value) | |
return level_to_args | |
def cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value): | |
def level_to_args(level): | |
level = int((level / MAX_LEVEL) * cutout_const) | |
return (level, replace_value) | |
return level_to_args | |
def solarize_level_to_args(MAX_LEVEL): | |
def level_to_args(level): | |
level = int((level / MAX_LEVEL) * 256) | |
return (level,) | |
return level_to_args | |
def none_level_to_args(level): | |
return () | |
def posterize_level_to_args(MAX_LEVEL): | |
def level_to_args(level): | |
level = int((level / MAX_LEVEL) * 4) | |
return (level,) | |
return level_to_args | |
def rotate_level_to_args(MAX_LEVEL, replace_value): | |
def level_to_args(level): | |
level = (level / MAX_LEVEL) * 30 | |
if np.random.random() < 0.5: | |
level = -level | |
return (level, replace_value) | |
return level_to_args | |
func_dict = { | |
"Identity": identity_func, | |
"AutoContrast": autocontrast_func, | |
"Equalize": equalize_func, | |
"Rotate": rotate_func, | |
"Solarize": solarize_func, | |
"Color": color_func, | |
"Contrast": contrast_func, | |
"Brightness": brightness_func, | |
"Sharpness": sharpness_func, | |
"ShearX": shear_x_func, | |
"TranslateX": translate_x_func, | |
"TranslateY": translate_y_func, | |
"Posterize": posterize_func, | |
"ShearY": shear_y_func, | |
} | |
translate_const = 10 | |
MAX_LEVEL = 10 | |
replace_value = (128, 128, 128) | |
arg_dict = { | |
"Identity": none_level_to_args, | |
"AutoContrast": none_level_to_args, | |
"Equalize": none_level_to_args, | |
"Rotate": rotate_level_to_args(MAX_LEVEL, replace_value), | |
"Solarize": solarize_level_to_args(MAX_LEVEL), | |
"Color": enhance_level_to_args(MAX_LEVEL), | |
"Contrast": enhance_level_to_args(MAX_LEVEL), | |
"Brightness": enhance_level_to_args(MAX_LEVEL), | |
"Sharpness": enhance_level_to_args(MAX_LEVEL), | |
"ShearX": shear_level_to_args(MAX_LEVEL, replace_value), | |
"TranslateX": translate_level_to_args(translate_const, MAX_LEVEL, replace_value), | |
"TranslateY": translate_level_to_args(translate_const, MAX_LEVEL, replace_value), | |
"Posterize": posterize_level_to_args(MAX_LEVEL), | |
"ShearY": shear_level_to_args(MAX_LEVEL, replace_value), | |
} | |
class RandomAugment(object): | |
def __init__(self, N=2, M=10, isPIL=False, augs=[]): | |
self.N = N | |
self.M = M | |
self.isPIL = isPIL | |
if augs: | |
self.augs = augs | |
else: | |
self.augs = list(arg_dict.keys()) | |
def get_random_ops(self): | |
sampled_ops = np.random.choice(self.augs, self.N) | |
return [(op, 0.5, self.M) for op in sampled_ops] | |
def __call__(self, img): | |
if self.isPIL: | |
img = np.array(img) | |
ops = self.get_random_ops() | |
for name, prob, level in ops: | |
if np.random.random() > prob: | |
continue | |
args = arg_dict[name](level) | |
img = func_dict[name](img, *args) | |
return img | |
class VideoRandomAugment(object): | |
def __init__(self, N=2, M=10, p=0.0, tensor_in_tensor_out=True, augs=[]): | |
self.N = N | |
self.M = M | |
self.p = p | |
self.tensor_in_tensor_out = tensor_in_tensor_out | |
if augs: | |
self.augs = augs | |
else: | |
self.augs = list(arg_dict.keys()) | |
def get_random_ops(self): | |
sampled_ops = np.random.choice(self.augs, self.N, replace=False) | |
return [(op, self.M) for op in sampled_ops] | |
def __call__(self, frames): | |
assert ( | |
frames.shape[-1] == 3 | |
), "Expecting last dimension for 3-channels RGB (b, h, w, c)." | |
if self.tensor_in_tensor_out: | |
frames = frames.numpy().astype(np.uint8) | |
num_frames = frames.shape[0] | |
ops = num_frames * [self.get_random_ops()] | |
apply_or_not = num_frames * [np.random.random(size=self.N) > self.p] | |
frames = torch.stack( | |
list(map(self._aug, frames, ops, apply_or_not)), dim=0 | |
).float() | |
return frames | |
def _aug(self, img, ops, apply_or_not): | |
for i, (name, level) in enumerate(ops): | |
if not apply_or_not[i]: | |
continue | |
args = arg_dict[name](level) | |
img = func_dict[name](img, *args) | |
return torch.from_numpy(img) | |
if __name__ == "__main__": | |
a = RandomAugment() | |
img = np.random.randn(32, 32, 3) | |
a(img) | |