|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
This code is refer from: |
|
https://github.com/RubanSeven/Text-Image-Augmentation-python/blob/master/augment.py |
|
""" |
|
|
|
import numpy as np |
|
from .warp_mls import WarpMLS |
|
|
|
|
|
def tia_distort(src, segment=4): |
|
img_h, img_w = src.shape[:2] |
|
|
|
cut = img_w // segment |
|
thresh = cut // 3 |
|
|
|
src_pts = list() |
|
dst_pts = list() |
|
|
|
src_pts.append([0, 0]) |
|
src_pts.append([img_w, 0]) |
|
src_pts.append([img_w, img_h]) |
|
src_pts.append([0, img_h]) |
|
|
|
dst_pts.append([np.random.randint(thresh), np.random.randint(thresh)]) |
|
dst_pts.append( |
|
[img_w - np.random.randint(thresh), np.random.randint(thresh)]) |
|
dst_pts.append( |
|
[img_w - np.random.randint(thresh), img_h - np.random.randint(thresh)]) |
|
dst_pts.append( |
|
[np.random.randint(thresh), img_h - np.random.randint(thresh)]) |
|
|
|
half_thresh = thresh * 0.5 |
|
|
|
for cut_idx in np.arange(1, segment, 1): |
|
src_pts.append([cut * cut_idx, 0]) |
|
src_pts.append([cut * cut_idx, img_h]) |
|
dst_pts.append([ |
|
cut * cut_idx + np.random.randint(thresh) - half_thresh, |
|
np.random.randint(thresh) - half_thresh |
|
]) |
|
dst_pts.append([ |
|
cut * cut_idx + np.random.randint(thresh) - half_thresh, |
|
img_h + np.random.randint(thresh) - half_thresh |
|
]) |
|
|
|
trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h) |
|
dst = trans.generate() |
|
|
|
return dst |
|
|
|
|
|
def tia_stretch(src, segment=4): |
|
img_h, img_w = src.shape[:2] |
|
|
|
cut = img_w // segment |
|
thresh = cut * 4 // 5 |
|
|
|
src_pts = list() |
|
dst_pts = list() |
|
|
|
src_pts.append([0, 0]) |
|
src_pts.append([img_w, 0]) |
|
src_pts.append([img_w, img_h]) |
|
src_pts.append([0, img_h]) |
|
|
|
dst_pts.append([0, 0]) |
|
dst_pts.append([img_w, 0]) |
|
dst_pts.append([img_w, img_h]) |
|
dst_pts.append([0, img_h]) |
|
|
|
half_thresh = thresh * 0.5 |
|
|
|
for cut_idx in np.arange(1, segment, 1): |
|
move = np.random.randint(thresh) - half_thresh |
|
src_pts.append([cut * cut_idx, 0]) |
|
src_pts.append([cut * cut_idx, img_h]) |
|
dst_pts.append([cut * cut_idx + move, 0]) |
|
dst_pts.append([cut * cut_idx + move, img_h]) |
|
|
|
trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h) |
|
dst = trans.generate() |
|
|
|
return dst |
|
|
|
|
|
def tia_perspective(src): |
|
img_h, img_w = src.shape[:2] |
|
|
|
thresh = img_h // 2 |
|
|
|
src_pts = list() |
|
dst_pts = list() |
|
|
|
src_pts.append([0, 0]) |
|
src_pts.append([img_w, 0]) |
|
src_pts.append([img_w, img_h]) |
|
src_pts.append([0, img_h]) |
|
|
|
dst_pts.append([0, np.random.randint(thresh)]) |
|
dst_pts.append([img_w, np.random.randint(thresh)]) |
|
dst_pts.append([img_w, img_h - np.random.randint(thresh)]) |
|
dst_pts.append([0, img_h - np.random.randint(thresh)]) |
|
|
|
trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h) |
|
dst = trans.generate() |
|
|
|
return dst |