Spaces:
Build error
Build error
File size: 3,190 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 |
import cv2
import numpy as np
from PIL import Image, ImageOps, ImageDraw
'''
PIL resize (W,H)
Torch resize is (H,W)
'''
class VGrid:
def __init__(self):
pass
def __call__(self, img, copy=True, max_width=4, mag=-1, prob=1.):
if np.random.uniform(0,1) > prob:
return img
if copy:
img = img.copy()
W, H = img.size
if mag<0 or mag>max_width:
line_width = np.random.randint(1, max_width)
image_stripe = np.random.randint(1, max_width)
else:
line_width = 1
image_stripe = 3 - mag
n_lines = W // (line_width + image_stripe) + 1
draw = ImageDraw.Draw(img)
for i in range(1, n_lines):
x = image_stripe*i + line_width*(i-1)
draw.line([(x,0), (x,H)], width=line_width, fill='black')
return img
class HGrid:
def __init__(self):
pass
def __call__(self, img, copy=True, max_width=4, mag=-1, prob=1.):
if np.random.uniform(0,1) > prob:
return img
if copy:
img = img.copy()
W, H = img.size
if mag<0 or mag>max_width:
line_width = np.random.randint(1, max_width)
image_stripe = np.random.randint(1, max_width)
else:
line_width = 1
image_stripe = 3 - mag
n_lines = H // (line_width + image_stripe) + 1
draw = ImageDraw.Draw(img)
for i in range(1, n_lines):
y = image_stripe*i + line_width*(i-1)
draw.line([(0,y), (W, y)], width=line_width, fill='black')
return img
class Grid:
def __init__(self):
pass
def __call__(self, img, mag=-1, prob=1.):
if np.random.uniform(0,1) > prob:
return img
img = VGrid()(img, copy=True, mag=mag)
img = HGrid()(img, copy=False, mag=mag)
return img
class RectGrid:
def __init__(self):
pass
def __call__(self, img, isellipse=False, mag=-1, prob=1.):
if np.random.uniform(0,1) > prob:
return img
img = img.copy()
W, H = img.size
line_width = 1
image_stripe = 3 - mag #np.random.randint(2, 6)
offset = 4 if isellipse else 1
n_lines = ((H//2) // (line_width + image_stripe)) + offset
draw = ImageDraw.Draw(img)
x_center = W // 2
y_center = H // 2
for i in range(1, n_lines):
dx = image_stripe*i + line_width*(i-1)
dy = image_stripe*i + line_width*(i-1)
x1 = x_center - (dx * W//H)
y1 = y_center - dy
x2 = x_center + (dx * W/H)
y2 = y_center + dy
if isellipse:
draw.ellipse([(x1,y1), (x2, y2)], width=line_width, outline='black')
else:
draw.rectangle([(x1,y1), (x2, y2)], width=line_width, outline='black')
return img
class EllipseGrid:
def __init__(self):
pass
def __call__(self, img, mag=-1, prob=1.):
if np.random.uniform(0,1) > prob:
return img
img = RectGrid()(img, isellipse=True, mag=mag, prob=prob)
return img
|