def gen_wm(RGB):
generator = GeneratorFromStrings(
strings=['eccv'],
count=1, # 五种字体
fonts=['fonts/Impact.ttf'], # default: []
language='en',
size=100, # 32
skewing_angle=10,
random_skew=False,
blur=0,
random_blur=False,
# gaussian noise (0), plain white (1), quasicrystal (2) or picture (3)
background_type=1,
distorsion_type=0, # None(0), Sine wave(1),Cosine wave(2),Random(3)
distorsion_orientation=0,
is_handwritten=False,
width=-1,
alignment=1,
text_color=RGB2Hex(RGB),
orientation=0,
space_width=1.0,
character_spacing=0,
margins=(0, 0, 0, 0),
fit=True,
)
img_list = [img for img, _ in generator]
return img_list[0]
def test_generator_from_strings(self):
generator = GeneratorFromStrings(["TEST TEST TEST"])
i = 0
while i < 100:
img, lbl = next(generator)
self.assertTrue(img.size[1] == 32, "Shape is not right")
i += 1
from trdg.generators import (
GeneratorFromDict,
GeneratorFromRandom,
GeneratorFromStrings,
GeneratorFromWikipedia,
)
generator = GeneratorFromStrings(
['Test1', 'Test2', 'Test3'],
blur=2,
random_blur=True
)
for img, lbl in generator:
# Do something with the pillow images here.
print(lbl)
def test_generator_from_strings_stops(self):
generator = GeneratorFromStrings(["TEST TEST TEST"], count=1)
next(generator)
self.assertRaises(StopIteration, generator.next)
from trdg.generators import GeneratorFromStrings
from PIL import ImageFont
from io import BytesIO
import os
font_dirs = [
os.path.join("mm_fonts", font_dir) for font_dir in os.listdir("mm_fonts")
]
all_fonts = []
for font_dir in font_dirs:
if os.path.isdir(font_dir):
for font_file in os.listdir(font_dir):
font_file = os.path.join(font_dir, font_file)
all_fonts.append(font_file)
print("All font dirs", all_fonts[0])
generator = GeneratorFromStrings(['ပြင်ပလင့်ခ်များငယ်ဘဝ'],
count=50,
fonts=all_fonts,
language="mm",
blur=2,
random_blur=True)
for img, lbl in generator:
img.save(os.path.join("test_images", lbl), format="png")
GeneratorFromRandom,
GeneratorFromStrings,
GeneratorFromWikipedia,
)
string = open("str/final.txt", "r+").read().split("\n")
random.shuffle(string)
# The generators use the same arguments as the CLI, only as parameters
generator = GeneratorFromStrings(
string,
blur=1,
count=30,
random_blur=True,
background_type=3,
random_skew=False,
space_width=1,
size=45,
character_spacing=1,
fit=True,
alignment=3,
margins=(5, 5, 5, 5),
#text_color="#a07081",
fonts=fonts,
image_dir="./background",
)
d = 0
for img, lbl in tqdm.tqdm(generator):
img.save("data/" + str(d) + ".jpg")
f = open("data/" + str(d) + ".txt", "w+")
f.write(lbl)
f.close()
def gen_box(dataloader, rects, title, count=40, contents_img_path="layout_process/contents_img"):
text_list, palette = dataloader.use_set()
palette = copy.deepcopy(palette)
sentences = []
# 제목 먼저 체크
if(title is not None):
# title이 만들어졌다면
start_x = title['x1']
start_y = title['y1']
end_x = title['x2']
end_y = title['y2']
word = []
sen_text = str()
while(start_y <= end_y):
if dataloader.gen_mode == True:
default_space=5
generator = GeneratorFromStrings(
[random.choice(text_list).replace('\n','') for _ in range(count)],
blur=0,
count=count,
background_type=0,
margins=(5, 5, 5, 5),
language='ko',
fonts=['../asset/fonts/ko/NanumBarunGothicUltraLight.ttf'],
random_blur=False
)
for text_img, text in generator:
if (end_x - start_x) >= text_img.width:
palette.paste(text_img, (start_x, start_y))
start_x += text_img.width + default_space
sen_text += text
word_dict = {
'box' : {
'x1' : start_x - text_img.width,
'x2' : start_x,
'y1' : start_y,
'y2' : start_y + 32,
},
'word':text
}
word.append(word_dict)
else:
for _ in range(len(text_list)):
text = random.choice(text_list)
if (end_x - start_x) >= text['width']:
text_img = Image.open(text['path'])
palette.paste(text_img, (start_x, start_y))
start_x += text['width']
sen_text += text['text']
word_dict = {
'box':{
'x1':start_x - text['width'],
'x2':start_x,
'y1':start_y,
'y2':start_y + 32,
},
'word':text['text']
}
word.append(word_dict)
sen_dict = {
'box':{
'x1' : title['x1'],
'x2' : start_x,
'y1' : start_y,
'y2' : start_y + 32
},
'sentence' : sen_text,
'words' : word
}
start_x = title['x1']
#need to specify
start_y += 32
sentences.append(sen_dict)
#add sentence box
sen_y = start_y
for layout in rects:
start_x = layout['x1']
start_y = layout['y1']
end_x = layout['x2']
end_y = layout['y2']
# 문단에 글 대신 이미지도 넣기
isImg = random.randint(6)
if(isImg == 3):
img_list = [files for files in os.listdir(contents_img_path) if ".png" in files] # 사진 파일들만 모으기
img_name = random.choice(img_list)
img_path = os.path.join(contents_img_path, img_name)
contents_img = Image.open(img_path) # 랜덤 이미지 오픈
contents_img = contents_img.resize((end_x - start_x, end_y - start_y)) # 문단 크기로 이미지 리사이즈
palette.paste(contents_img, (start_x, start_y))
continue # 이미지 넣었으면 그 문단은 패스
# 여기서부터 텍스트 이미지
trial = 0
while True:
trial += 1
sen_x = start_x
sen_text = str()
word = []
#print(f'gen_trial : {trial}')
#if gen_mode is true
if dataloader.gen_mode == True:
word = []
default_space=5
generator = GeneratorFromStrings(
[random.choice(text_list).replace('\n','') for _ in range(count)],
blur=0,
count=count,
background_type=0,
margins=(5, 5, 5, 5),
language='ko',
fonts=['../asset/fonts/ko/NanumBarunGothicUltraLight.ttf'],
random_blur=False
)
for text_img, text in generator:
if (end_x - start_x) >= text_img.width:
palette.paste(text_img, (start_x, start_y))
start_x += text_img.width + default_space
sen_text += text
word_dict = {
'box' : {
'x1' : start_x - text_img.width,
'x2' : start_x,
'y1' : start_y,
'y2' : start_y + 32,
},
'word':text
}
word.append(word_dict)
# if gen_mode is false use default dict shape dataset
else:
for _ in range(len(text_list)):
text = random.choice(text_list)
if (end_x - start_x) >= text['width']:
text_img = Image.open(text['path'])
palette.paste(text_img, (start_x, start_y))
start_x += text['width']
sen_text += text['text']
word_dict = {
'box':{
'x1':start_x - text['width'],
'x2':start_x,
'y1':start_y,
'y2':start_y + 32,
},
'word':text['text']
}
word.append(word_dict)
sen_dict = {
'box':{
'x1' : layout['x1'],
'x2' : start_x,
'y1' : start_y,
'y2' : start_y + 32
},
'sentence' : sen_text,
'words' : word
}
start_x = layout['x1']
#need to specify
start_y += 32
sentences.append(sen_dict)
#add sentence box
sen_y = start_y
if start_y > end_y:
break
return sentences, palette
import glob
import random
import tqdm
import string
fonts = glob.glob("font/name.otf")
# string = open('data_gen/name/name.txt').read().splitlines()
generator = GeneratorFromStrings(
['X.Liêm Thuỷ, H.Na Rì, T.Bắc Kạn'],
blur=1,
count=20,
random_blur=True,
language='en',
background_type=3,
skewing_angle=0,
random_skew=False,
space_width=1,
size=45,
character_spacing=1,
fit=True,
alignment=3,
margins=(5, 5, 5, 5),
fonts=fonts,
image_dir="./image_long",
)
d = 0
for img, lbl in tqdm.tqdm(generator):
img.save("data/" + str(d) + ".jpg")
with open('data/' + str(d) +'.txt', 'w') as f:
f.write(str(lbl))
d += 1
def generate(save_img_dir,
save_txt=None,
remake=False,
im_num=10,
rotation=[0, 180]):
# The generators use the same arguments as the CLI, only as parameters
generator = GeneratorFromStrings(
corpus,
blur=2,
size=64,
language="msttcorefonts",
skewing_angle=0,
# random_skew=True,
margins=(0, 0, 10, 0), # up,left,bottom, right
random_blur=True)
save_img_dir = os.path.join(save_img_dir, 'synthetic')
if remake:
if os.path.exists(save_img_dir):
shutil.rmtree(save_img_dir)
for r in rotation:
if not os.path.exists(os.path.join(save_img_dir, str(r))):
os.makedirs(os.path.join(save_img_dir, str(r)))
c = 0
labels = []
for img, lbl in generator:
if c > im_num: break
if c % 1000 == 0:
print("generated {} images".format(c))
im_name = str(c) + '.png'
im_path = os.path.join(save_img_dir, str(generator.skewing_angle),
im_name)
img.save(im_path)
labels.append(im_path + ',' +
str(1 if not generator.skewing_angle else 0) + ',' +
str(0 if not generator.skewing_angle else 1))
c += 1
skewing_angle = random.choice(rotation)
if skewing_angle == 0:
total_hz_mg = randint(1, 30)
total_vz_mg = randint(1, 15)
left_mg = randint(1, total_hz_mg)
right_mg = total_hz_mg - left_mg
up_mg = randint(1, total_vz_mg)
down_mg = total_vz_mg - up_mg + 10
generator.margins = (up_mg, left_mg, down_mg, right_mg
) # up, left, down, right
else:
total_hz_mg = randint(1, 30)
total_vz_mg = randint(1, 15)
left_mg = randint(1, total_hz_mg)
right_mg = total_hz_mg - left_mg
down_mg = randint(1, total_vz_mg)
up_mg = total_vz_mg - down_mg + 10
generator.margins = (up_mg, left_mg, down_mg, right_mg
) # up, left, down, right
generator.skewing_angle = skewing_angle
if save_txt:
labels_txt = '\n'.join(labels)
if not os.path.exists(os.path.join(save_img_dir, save_txt)):
with open(os.path.join(save_img_dir, save_txt), 'w') as f:
f.write(labels_txt)
else:
with open(os.path.join(save_img_dir, save_txt), 'a+') as f:
f.seek(0, 1)
f.seek(f.tell() - 1, 0)
last_char = f.read()
if last_char == '\n':
f.write(labels_txt)
else:
f.write('\n')
f.write(labels_txt)
print('generate synthetic images done!')
def main():
fonts = multiline_load_fonts(font_dir, font)
lang_dict = []
if os.path.isfile(dict_dir):
with open(dict_dir, "r", encoding="utf8", errors="ignore") as d:
lang_dict = [l for l in d.read().splitlines() if len(l) > 0]
else:
sys.exit("Cannot open dict")
strings1 = create_strings_from_dict(length[0], random_length[0], count,
lang_dict)
strings2 = create_strings_from_dict(length[1], random_length[1], count,
lang_dict)
for i in range(count):
generator1 = GeneratorFromStrings(strings=rnd.choices(strings1),
count=len(strings1),
size=font_size[0],
bounding_box=True,
fonts=fonts,
background_type=1,
text_color="#FFFFFF")
img1, bbox1, lb1 = unpack_generator(generator1)
generator2 = GeneratorFromStrings(strings=rnd.choices(strings2),
size=font_size[1],
count=len(strings2),
bounding_box=True,
fonts=fonts,
background_type=1,
text_color="#FFFFFF")
img2, bbox2, lb2 = unpack_generator(generator2)
#first_line = copy.deepcopy(rnd.choice(generator1_list))
#second_line = copy.deepcopy(rnd.choice(generator2_list))
#img1, bbox1, lb1 = first_line
#img2, bbox2, lb2 = second_line
img = generate_background_img(background_img_dir, image_width,
image_height)
img_draw = ImageDraw.Draw(img)
x = rnd.randint(0, img.width - max(img1.width, img2.width))
y = rnd.randint(0, img.height - img1.height - img2.height)
###### paste first line ######
img.paste(img1, (x, y), img1)
bbox1 = paste_bbox(bbox1, x, y)
'''
for box in bbox1:
img_draw.polygon(box, outline = (255, 0, 0), )
'''
###### paste first line ######
###### paste second line ######
img.paste(img2, (x, y + img1.height + line_space), img2)
bbox2 = paste_bbox(bbox2, x, y + img1.height + line_space)
'''
for box in bbox2:
img_draw.polygon(box, outline = (0, 255, 0), )
'''
###### paste second line ######
img = img.convert("RGB")
img.save(out_dir + lb1 + lb2 + ".jpg")
with open(os.path.join(out_dir, lb1 + lb2 + '.txt'), mode='w') as f:
for box, p in zip(bbox1 + bbox2, lb1 + lb2):
f.write(str(box) + p)
f.write('\r\n')
strings.append(current_string)
return strings
# The generators use the same arguments as the CLI, only as parameters
# strs = create_strings_randomly(100, True, True, True, lang='en')
strs = create_strings_from_dict(10,
allow_variable=True,
count=5000,
dict_path='dict.txt')
generator = GeneratorFromStrings(strings=strs,
blur=0,
random_blur=True,
language='cn',
skewing_angle=10,
random_skew=True,
distorsion_type=3,
margins=(0, 0, 0, 0),
size=64)
save_dir = r'/home/zj/桌面/ocr_simple_dataset'
img_folder = os.path.join(save_dir, 'val')
if os.path.exists(img_folder):
shutil.rmtree(img_folder, ignore_errors=True)
os.makedirs(img_folder, exist_ok=True)
pbar = tqdm(total=len(strs))
with open(os.path.join(save_dir, 'val.txt'), mode='w', encoding='utf8') as f:
i = 0
for img, text in generator:
img_path = os.path.join(img_folder, '{}.jpg'.format(i))