forked from Jayhello/python_utils
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathimg_augmentation.py
More file actions
171 lines (135 loc) · 5.05 KB
/
Copy pathimg_augmentation.py
File metadata and controls
171 lines (135 loc) · 5.05 KB
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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
# _*_ coding:utf-8 _*_
"""
Deep learning image augmentation
cited from https://scottontechnology.com/flip-image-opencv-python/
http://augmentor.readthedocs.io/en/master/userguide/mainfeatures.html
"""
import cv2
import glob
import random
import os
from multiprocessing import Pool as ProcessPool
from multiprocessing.dummy import Pool as ThreadPool
import Augmentor
import path_var
def img_flip():
# path = "F:/ad_samples/train_samples/ad_text_artifact/base_type/type_10.jpg"
path = "F:/ad_samples/download_sample/14/8DB54D749B1D4A2D5FD3441C681D9A2C522453AC_s.jpg"
img = cv2.imread(path)
horizontal_img = img.copy()
vertical_img = img.copy()
both_img = img.copy()
horizontal_img = cv2.flip(img, 0)
vertical_img = cv2.flip(img, 1)
both_img = cv2.flip(img, -1)
cv2.imshow("original img", img)
cv2.imshow("horizontal img", horizontal_img)
cv2.imshow("vertical img", vertical_img)
cv2.imshow("both flip", both_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
# g_dst_dir = 'F:/ad_samples/train_samples/artifact_ad_web/'
# g_dst_dir = 'F:/ad_samples/train_samples/artifact_ad_web_2_flip/'
# class Path(object):
# g_dst_dir = ''
# global g_dst_dir
# g_dst_dir = ''
def flip_img_save2dir(file):
img = cv2.imread(file)
dst_dir = path_var.g_dst_dir
h_img = img.copy()
v_img = img.copy()
b_img = img.copy()
h_img = cv2.flip(img, 0)
v_img = cv2.flip(img, 1)
b_img = cv2.flip(img, -1)
# file like F:/ad_samples/train_samples/ad_text_artifact/base_type/type_10.jpg
# get file name "type_10"
# type_10.jpg
base_name = os.path.basename(file)
# type_10
base_name = os.path.splitext(base_name)[0]
file_name = dst_dir + base_name + "_h" + ".jpg"
cv2.imwrite(file_name, h_img)
file_name = dst_dir + base_name + "_v" + ".jpg"
cv2.imwrite(file_name, v_img)
file_name = dst_dir + base_name + '_b' + ".jpg"
cv2.imwrite(file_name, b_img)
def do_all_flip(base_dir="F:/ad_samples/train_samples/ad_web_2/"):
"""
flip all the images in dir, and then save them
to another dir
:return:
"""
# get all files
# base_dir = "F:/ad_samples/train_samples/ad_text/"
# base_dir = "F:/ad_samples/download_sample/14/"
# base_dir = "F:/ad_samples/train_samples/ad_web/"
# files = glob.glob(base_dir + "*.jpg")
files = glob.glob(base_dir + "/*.png")
# like ['E:/img\\1.jpg', 'E:/img\\10.jpg']
# start 3 process
# pool = ProcessPool(3)
pool = ThreadPool(3)
rets = pool.map(flip_img_save2dir, files)
pool.close()
pool.join()
print 'all images accomplish flip and save to dir'
def flip_all_in_dir():
base_dir = 'F:/ad_samples/train_samples/others/'
sub_dir_lst = glob.glob(base_dir + "*")
# ['F:/dir1', 'F:/dir2']
# print sub_dir_lst
new_sub_dir = [os.path.join(base_dir, item + '_flip/') for item in os.listdir(base_dir)]
# ['F:/dir1_flip', 'F:/dir2_flip']
for dir_item, new_item in zip(sub_dir_lst[10:], new_sub_dir[10:]):
global g_dst_dir
if not os.path.exists(new_item):
os.makedirs(new_item)
# g_dst_dir = new_item
# Path.g_dst_dir = new_item
path_var.g_dst_dir = new_item
print 'flip %s, flip dir %s' % (dir_item, new_item)
do_all_flip(base_dir=dir_item)
# print os.listdir(base_dir)
# ['dir1', 'dir2']
# print glob.glob(base_dir + "*")
# ['F:/dir1', 'F:/dir2']
def augmentation():
path = 'F:/augment'
# path = 'F:/ad_samples/train_samples/ad_text'
# output_path = 'F:/ad_samples/train_samples/ad_text_artifact/augmentation'
output_path = 'output'
# path = 'F:/ad_samples/train_samples/artifact_ad_web/artifact_ad_web_14_flip/'
# path = 'F:/ad_samples/train_samples/ad_text_2/'
# path = 'F:/ad_samples/new_train_sample/'
# output_path = 'F:/ad_samples/train_samples/artifact_ad_web/artifact_ad_web_14_aug/'
# output_path = 'F:/ad_samples/train_samples/ad_text_rotate90/'
# output_path = 'F:/ad_samples/new_train_sample_artifact/'
# output_path = 'F:/ad_samples/train_samples/ad_text_2_output/'
p = Augmentor.Pipeline(path, output_directory=output_path)
p.zoom(probability=0.1, min_factor=1.1, max_factor=1.3)
p.flip_left_right(probability=0.1)
p.rotate(probability=0.2, max_left_rotation=15, max_right_rotation=16)
p.shear(probability=0.2, max_shear_left=10, max_shear_right=10)
p.skew(probability=0.1, magnitude=0.6)
p.skew_tilt(probability=0.2, magnitude=0.6)
p.random_distortion(probability=0.3, grid_height=4, grid_width=4, magnitude=4)
# p.random_distortion(probability=0.2, grid_height=4, grid_width=4, magnitude=4)
# p.rotate90(probability=1)
# SIZE = 4164 * 4
SIZE = 5 * 4
# SIZE = 2799 * 10
# SIZE = 2799 * 1
# SIZE = 4148 * 3
# SIZE = 68 * 5
p.sample(SIZE)
if __name__ == '__main__':
# img_flip()
# flip_all_in_dir()
# do_all_flip()
augmentation()
# test single image flip and save
# file = 'F:/ad_samples/train_samples/ad_text_artifact/base_type/type_10.jpg'
# flip_img_save2dir(file=file)
pass