Spaces:
Build error
Build error
| import cv2 | |
| import numpy as np | |
| import skimage as sk | |
| from PIL import Image, ImageOps | |
| from io import BytesIO | |
| from skimage import color | |
| ''' | |
| PIL resize (W,H) | |
| cv2 image is BGR | |
| PIL image is RGB | |
| ''' | |
| class Contrast: | |
| def __init__(self): | |
| pass | |
| def __call__(self, img, mag=-1, prob=1.): | |
| if np.random.uniform(0,1) > prob: | |
| return img | |
| #c = [0.4, .3, .2, .1, .05] | |
| c = [0.4, .3, .2] | |
| if mag<0 or mag>=len(c): | |
| index = np.random.randint(0, len(c)) | |
| else: | |
| index = mag | |
| c = c[index] | |
| img = np.array(img) / 255. | |
| means = np.mean(img, axis=(0, 1), keepdims=True) | |
| img = np.clip((img - means) * c + means, 0, 1) * 255 | |
| return Image.fromarray(img.astype(np.uint8)) | |
| class Brightness: | |
| def __init__(self): | |
| pass | |
| def __call__(self, img, mag=-1, prob=1.): | |
| if np.random.uniform(0,1) > prob: | |
| return img | |
| #W, H = img.size | |
| #c = [.1, .2, .3, .4, .5] | |
| c = [.1, .2, .3] | |
| if mag<0 or mag>=len(c): | |
| index = np.random.randint(0, len(c)) | |
| else: | |
| index = mag | |
| c = c[index] | |
| n_channels = len(img.getbands()) | |
| isgray = n_channels == 1 | |
| img = np.array(img) / 255. | |
| if isgray: | |
| img = np.expand_dims(img, axis=2) | |
| img = np.repeat(img, 3, axis=2) | |
| img = sk.color.rgb2hsv(img) | |
| img[:, :, 2] = np.clip(img[:, :, 2] + c, 0, 1) | |
| img = sk.color.hsv2rgb(img) | |
| #if isgray: | |
| # img = img[:,:,0] | |
| # img = np.squeeze(img) | |
| img = np.clip(img, 0, 1) * 255 | |
| img = Image.fromarray(img.astype(np.uint8)) | |
| if isgray: | |
| img = ImageOps.grayscale(img) | |
| return img | |
| #if isgray: | |
| #if isgray: | |
| # img = color.rgb2gray(img) | |
| #return Image.fromarray(img.astype(np.uint8)) | |
| class JpegCompression: | |
| def __init__(self): | |
| pass | |
| def __call__(self, img, mag=-1, prob=1.): | |
| if np.random.uniform(0,1) > prob: | |
| return img | |
| #c = [25, 18, 15, 10, 7] | |
| c = [25, 18, 15] | |
| if mag<0 or mag>=len(c): | |
| index = np.random.randint(0, len(c)) | |
| else: | |
| index = mag | |
| c = c[index] | |
| output = BytesIO() | |
| img.save(output, 'JPEG', quality=c) | |
| return Image.open(output) | |
| class Pixelate: | |
| def __init__(self): | |
| pass | |
| def __call__(self, img, mag=-1, prob=1.): | |
| if np.random.uniform(0,1) > prob: | |
| return img | |
| W, H = img.size | |
| #c = [0.6, 0.5, 0.4, 0.3, 0.25] | |
| c = [0.6, 0.5, 0.4] | |
| if mag<0 or mag>=len(c): | |
| index = np.random.randint(0, len(c)) | |
| else: | |
| index = mag | |
| c = c[index] | |
| img = img.resize((int(W* c), int(H * c)), Image.BOX) | |
| return img.resize((W, H), Image.BOX) | |