Spaces:
Build error
Build error
File size: 2,904 Bytes
d61b9c7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
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)
|