Skip to content

Commit 40da5fe

Browse files
n0obcodern0obcoder
authored
Update generate_synthetic_dataset.py
1 parent a3bc725 commit 40da5fe

File tree

1 file changed

+116
-133
lines changed

1 file changed

+116
-133
lines changed

‎generate_synthetic_dataset.py

+116-133
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ def degrade_quality(img):
4545
if not os.path.exists(dir_path):
4646
os.mkdir(dir_path)
4747

48-
f = open("shitty_text.txt", "r")
48+
f = open("shitty_text.txt", encoding='utf-8', mode="r")
4949
text = f.read()
5050
f.close()
5151
lines_list = str.split(text, '\n')
@@ -78,138 +78,121 @@ def degrade_quality(img):
7878
word_start_x = 5
7979
word_end_y = 5
8080
for i in tqdm(range(num_imgs)):
81-
# try:
82-
# make a blank image
83-
img = np.ones((h, w), dtype = np.uint8)*255
84-
85-
# set random parameters
86-
font = font_list[np.random.randint(len(font_list))]
87-
bottomLeftCornerOfText = (np.random.randint(word_start_x, int(img.shape[1]/3)), np.random.randint(int(img.shape[0]/2), int(img.shape[0]) - word_end_y))
88-
fontColor = np.random.randint(0,30)# (np.random.randint(0,30),np.random.randint(0,30),np.random.randint(0,30))
89-
fontScale = np.random.randint(2200,3000)/1000
90-
lineType = np.random.randint(1,3)
91-
92-
# text to be printed on the blank image
93-
num_words = np.random.randint(1,8)
94-
print_text = ''
95-
for _ in range(num_words):
96-
print_text += str.split(words_list[word_count])[0] + ' '
97-
word_count += 1
98-
print_text = print_text[:-1] # to get rif of the last space
99-
100-
# writing the text on the image
101-
cv2.putText(img, print_text.upper(), bottomLeftCornerOfText, font, fontScale, fontColor, lineType)
102-
noisy_img = img.copy()
103-
#################################################
104-
# add horizontal line at the bottom of the text
105-
106-
black_coords = np.where(noisy_img == fontColor)
107-
108-
ymin = np.min(black_coords[0])
109-
ymax = np.max(black_coords[0])
110-
xmin = np.min(black_coords[1])
111-
xmax = np.max(black_coords[1])
112-
113-
h_start_x = 0 #np.random.randint(0,xmin)
114-
h_end_x = np.random.randint(int(img.shape[1]*0.8), img.shape[1])
115-
h_length = h_end_x - h_start_x + 1
116-
num_h_lines = np.random.randint(10,30)
117-
h_lines = []
118-
h_start_temp = h_start_x
119-
# print('num_h_lines: ', num_h_lines)
120-
121-
next_line = True
122-
num_line = 0
123-
while (next_line) and (num_line < num_h_lines):
124-
if h_start_temp < h_end_x:
125-
# print('h_start_temp: ', h_start_temp)
126-
h_end_temp = np.random.randint(h_start_temp + 1, h_end_x + 1)
127-
# print('h_end_temp : ', h_end_temp)
128-
if h_end_temp < h_end_x:
129-
# print('0')
130-
h_lines.append([h_start_temp, h_end_temp])
131-
h_start_temp = h_end_temp + 1
132-
num_line += 1
133-
else:
134-
h_lines.append([h_start_temp, h_end_x])
135-
num_line += 1
136-
# print('1')
137-
next_line = False
81+
# make a blank image
82+
img = np.ones((h, w), dtype = np.uint8)*255
83+
84+
# set random parameters
85+
font = font_list[np.random.randint(len(font_list))]
86+
bottomLeftCornerOfText = (np.random.randint(word_start_x, int(img.shape[1]/3)), np.random.randint(int(img.shape[0]/2), int(img.shape[0]) - word_end_y))
87+
fontColor = np.random.randint(0,30)# (np.random.randint(0,30),np.random.randint(0,30),np.random.randint(0,30))
88+
fontScale = np.random.randint(2200,3000)/1000
89+
lineType = np.random.randint(1,3)
90+
91+
# text to be printed on the blank image
92+
num_words = np.random.randint(1,8)
93+
print_text = ''
94+
for _ in range(num_words):
95+
print_text += str.split(words_list[word_count])[0] + ' '
96+
word_count += 1
97+
print_text = print_text[:-1] # to get rif of the last space
98+
99+
# writing the text on the image
100+
cv2.putText(img, print_text.upper(), bottomLeftCornerOfText, font, fontScale, fontColor, lineType)
101+
noisy_img = img.copy()
102+
#################################################
103+
# add horizontal line at the bottom of the text
104+
105+
black_coords = np.where(noisy_img == fontColor)
106+
107+
ymin = np.min(black_coords[0])
108+
ymax = np.max(black_coords[0])
109+
xmin = np.min(black_coords[1])
110+
xmax = np.max(black_coords[1])
111+
112+
h_start_x = 0 #np.random.randint(0,xmin)
113+
h_end_x = np.random.randint(int(img.shape[1]*0.8), img.shape[1])
114+
h_length = h_end_x - h_start_x + 1
115+
num_h_lines = np.random.randint(10,30)
116+
h_lines = []
117+
h_start_temp = h_start_x
118+
# print('num_h_lines: ', num_h_lines)
119+
120+
next_line = True
121+
num_line = 0
122+
while (next_line) and (num_line < num_h_lines):
123+
if h_start_temp < h_end_x:
124+
h_end_temp = np.random.randint(h_start_temp + 1, h_end_x + 1)
125+
# print('h_end_temp : ', h_end_temp)
126+
if h_end_temp < h_end_x:
127+
h_lines.append([h_start_temp, h_end_temp])
128+
h_start_temp = h_end_temp + 1
129+
num_line += 1
138130
else:
139-
# print('2')
131+
h_lines.append([h_start_temp, h_end_x])
132+
num_line += 1
140133
next_line = False
141-
142-
for h_line in h_lines:
143-
# print('h_line: ', h_line)
144-
col = np.random.choice(['black', 'white'], p = [0.65, 0.35])
145-
if col == 'black':
146-
# cv2.line(noisy_img, (h_line[0], ymax), (h_line[1], ymax), np.random.randint(0, 30), np.random.randint(3, 5))
147-
148-
149-
x_points = list(range(h_line[0], h_line[1] + 1))
150-
x_points_black_prob = np.random.choice([0,1], size = len(x_points), p = [0.2, 0.8])
151-
152-
for idx, x in enumerate(x_points):
153-
if x_points_black_prob[idx]:
154-
noisy_img[ ymax - np.random.randint(4): ymax + np.random.randint(4), x] = np.random.randint(0,30)
155-
156-
157-
158-
# print('---')
159-
#################################################
160-
# adding vertical lines
161-
vertical_bool = {'left': np.random.choice([0,1], p =[0.2, 0.8]), 'right': np.random.choice([0,1])} # [1 or 0, 1 or 0] whether to make vertical left line on left and right side of the image
162-
for left_right, bool_ in vertical_bool.items():
163-
# print(left_right, bool_)
164-
if bool_:
165-
if left_right == 'left':
166-
v_start_x = np.random.randint(5, int(noisy_img.shape[1]*0.06))
167-
else:
168-
v_start_x = np.random.randint(int(noisy_img.shape[1]*0.95), noisy_img.shape[1] - 5)
169-
170-
v_start_y = np.random.randint(0, int(noisy_img.shape[0]*0.06))
171-
v_end_y = np.random.randint(int(noisy_img.shape[0]*0.95), noisy_img.shape[0])
172-
173-
y_points = list(range(v_start_y, v_end_y + 1))
174-
y_points_black_prob = np.random.choice([0,1], size = len(y_points), p = [0.2, 0.8])
175-
176-
for idx, y in enumerate(y_points):
177-
if y_points_black_prob[idx]:
178-
noisy_img[y, v_start_x - np.random.randint(4): v_start_x + np.random.randint(4)] = np.random.randint(0,30)
179-
180-
# '''
181-
# erode the image
182-
kernel = np.ones((3,3),np.uint8)
183-
erosion_iteration = np.random.randint(1,3)
184-
dilate_iteration = np.random.randint(0,2)
185-
img = cv2.erode(img,kernel,iterations = erosion_iteration)
186-
noisy_img = cv2.erode(noisy_img,kernel,iterations = erosion_iteration)
187-
img = cv2.dilate(img,kernel,iterations = dilate_iteration)
188-
noisy_img = cv2.dilate(noisy_img,kernel,iterations = dilate_iteration)
189-
# '''
190-
191-
img = degrade_quality(img)
192-
noisy_img = degrade_quality(noisy_img)
193-
194-
debug_img = np.ones((2*h, w), dtype = np.uint8)*255
195-
debug_img[0:h, :] = img
196-
debug_img[h:2*h, :] = noisy_img
197-
cv2.line(debug_img, (0, h), (debug_img.shape[1], h), 150, 5)
198-
199-
img = cv2.resize(img, (0,0), fx = 0.25, fy = 0.25)
200-
noisy_img = cv2.resize(noisy_img, (0,0), fx = 0.25, fy = 0.25)
201-
debug_img = cv2.resize(debug_img, (0,0), fx = 0.25, fy = 0.25)
202-
203-
if img_count <= train_num:
204-
cv2.imwrite(os.path.join(data_dir, 'train', 'imgs', '{}.jpg'.format(str(img_count).zfill(6))), img)
205-
cv2.imwrite(os.path.join(data_dir, 'train', 'noisy', '{}.jpg'.format(str(img_count).zfill(6))), noisy_img)
206-
cv2.imwrite(os.path.join(data_dir, 'train', 'debug', '{}.jpg'.format(str(img_count).zfill(6))), debug_img)
207134
else:
208-
cv2.imwrite(os.path.join(data_dir, 'val', 'imgs', '{}.jpg'.format(str(img_count).zfill(6))), img)
209-
cv2.imwrite(os.path.join(data_dir, 'val', 'noisy', '{}.jpg'.format(str(img_count).zfill(6))), noisy_img)
210-
cv2.imwrite(os.path.join(data_dir, 'val', 'debug', '{}.jpg'.format(str(img_count).zfill(6))), debug_img)
211-
212-
img_count += 1
213-
214-
# except:
215-
# pass
135+
next_line = False
136+
137+
for h_line in h_lines:
138+
col = np.random.choice(['black', 'white'], p = [0.65, 0.35])
139+
if col == 'black':
140+
x_points = list(range(h_line[0], h_line[1] + 1))
141+
x_points_black_prob = np.random.choice([0,1], size = len(x_points), p = [0.2, 0.8])
142+
143+
for idx, x in enumerate(x_points):
144+
if x_points_black_prob[idx]:
145+
noisy_img[ ymax - np.random.randint(4): ymax + np.random.randint(4), x] = np.random.randint(0,30)
146+
147+
# adding vertical lines
148+
vertical_bool = {'left': np.random.choice([0,1], p =[0.2, 0.8]), 'right': np.random.choice([0,1])} # [1 or 0, 1 or 0] whether to make vertical left line on left and right side of the image
149+
for left_right, bool_ in vertical_bool.items():
150+
if bool_:
151+
if left_right == 'left':
152+
v_start_x = np.random.randint(5, int(noisy_img.shape[1]*0.06))
153+
else:
154+
v_start_x = np.random.randint(int(noisy_img.shape[1]*0.95), noisy_img.shape[1] - 5)
155+
156+
v_start_y = np.random.randint(0, int(noisy_img.shape[0]*0.06))
157+
v_end_y = np.random.randint(int(noisy_img.shape[0]*0.95), noisy_img.shape[0])
158+
159+
y_points = list(range(v_start_y, v_end_y + 1))
160+
y_points_black_prob = np.random.choice([0,1], size = len(y_points), p = [0.2, 0.8])
161+
162+
for idx, y in enumerate(y_points):
163+
if y_points_black_prob[idx]:
164+
noisy_img[y, v_start_x - np.random.randint(4): v_start_x + np.random.randint(4)] = np.random.randint(0,30)
165+
166+
# '''
167+
# erode the image
168+
kernel = np.ones((3,3),np.uint8)
169+
erosion_iteration = np.random.randint(1,3)
170+
dilate_iteration = np.random.randint(0,2)
171+
img = cv2.erode(img,kernel,iterations = erosion_iteration)
172+
noisy_img = cv2.erode(noisy_img,kernel,iterations = erosion_iteration)
173+
img = cv2.dilate(img,kernel,iterations = dilate_iteration)
174+
noisy_img = cv2.dilate(noisy_img,kernel,iterations = dilate_iteration)
175+
# '''
176+
177+
img = degrade_quality(img)
178+
noisy_img = degrade_quality(noisy_img)
179+
180+
debug_img = np.ones((2*h, w), dtype = np.uint8)*255
181+
debug_img[0:h, :] = img
182+
debug_img[h:2*h, :] = noisy_img
183+
cv2.line(debug_img, (0, h), (debug_img.shape[1], h), 150, 5)
184+
185+
img = cv2.resize(img, (0,0), fx = 0.25, fy = 0.25)
186+
noisy_img = cv2.resize(noisy_img, (0,0), fx = 0.25, fy = 0.25)
187+
debug_img = cv2.resize(debug_img, (0,0), fx = 0.25, fy = 0.25)
188+
189+
if img_count <= train_num:
190+
cv2.imwrite(os.path.join(data_dir, 'train', 'imgs', '{}.jpg'.format(str(img_count).zfill(6))), img)
191+
cv2.imwrite(os.path.join(data_dir, 'train', 'noisy', '{}.jpg'.format(str(img_count).zfill(6))), noisy_img)
192+
cv2.imwrite(os.path.join(data_dir, 'train', 'debug', '{}.jpg'.format(str(img_count).zfill(6))), debug_img)
193+
else:
194+
cv2.imwrite(os.path.join(data_dir, 'val', 'imgs', '{}.jpg'.format(str(img_count).zfill(6))), img)
195+
cv2.imwrite(os.path.join(data_dir, 'val', 'noisy', '{}.jpg'.format(str(img_count).zfill(6))), noisy_img)
196+
cv2.imwrite(os.path.join(data_dir, 'val', 'debug', '{}.jpg'.format(str(img_count).zfill(6))), debug_img)
197+
198+
img_count += 1

0 commit comments

Comments
 (0)