108 lines
2.8 KiB
Python
108 lines
2.8 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: RubanSeven
|
|
# Reference: https://github.com/RubanSeven/Text-Image-Augmentation-python
|
|
|
|
# import cv2
|
|
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
|