def test_generate_data_with_white_background(self):
background_generator.plain_white(
64, 128).convert("RGB").save("tests/out/white_background.jpg")
self.assertTrue(
md5("tests/out/white_background.jpg") == md5(
"tests/expected_results/white_background.jpg"))
os.remove("tests/out/white_background.jpg")
def test_generate_data_with_white_background(self):
background_generator.plain_white(
64, 128).convert("RGB").save("tests/out/white_background.jpg")
self.assertTrue(
diff(
"tests/out/white_background.jpg",
"tests/expected_results/white_background.jpg",
delete_diff_file=True,
) < 0.01)
os.remove("tests/out/white_background.jpg")
def generate(
cls,
index,
text,
font,
out_dir,
size,
extension,
skewing_angle,
random_skew,
blur,
random_blur,
background_type,
distorsion_type,
distorsion_orientation,
is_handwritten,
name_format,
width,
alignment,
text_color,
orientation,
space_width,
character_spacing,
margins,
fit,
output_mask,
word_split,
image_dir,
):
image = None
margin_top, margin_left, margin_bottom, margin_right = margins
horizontal_margin = margin_left + margin_right
vertical_margin = margin_top + margin_bottom
##########################
# Create picture of text #
##########################
if is_handwritten:
if orientation == 1:
raise ValueError("Vertical handwritten text is unavailable")
image, mask = handwritten_text_generator.generate(text, text_color)
else:
image, mask = computer_text_generator.generate(
text,
font,
text_color,
size,
orientation,
space_width,
character_spacing,
fit,
word_split,
)
random_angle = rnd.randint(0 - skewing_angle, skewing_angle)
rotated_img = image.rotate(
skewing_angle if not random_skew else random_angle, expand=1)
rotated_mask = mask.rotate(
skewing_angle if not random_skew else random_angle, expand=1)
#############################
# Apply distorsion to image #
#############################
if distorsion_type == 0:
distorted_img = rotated_img # Mind = blown
distorted_mask = rotated_mask
elif distorsion_type == 1:
distorted_img, distorted_mask = distorsion_generator.sin(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
elif distorsion_type == 2:
distorted_img, distorted_mask = distorsion_generator.cos(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
else:
distorted_img, distorted_mask = distorsion_generator.random(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
##################################
# Resize image to desired format #
##################################
# Horizontal text
if orientation == 0:
new_width = int(
distorted_img.size[0] *
(float(size - vertical_margin) / float(distorted_img.size[1])))
resized_img = distorted_img.resize(
(new_width, size - vertical_margin), Image.ANTIALIAS)
resized_mask = distorted_mask.resize(
(new_width, size - vertical_margin), Image.NEAREST)
background_width = width if width > 0 else new_width + horizontal_margin
background_height = size
# Vertical text
elif orientation == 1:
new_height = int(
float(distorted_img.size[1]) *
(float(size - horizontal_margin) /
float(distorted_img.size[0])))
resized_img = distorted_img.resize(
(size - horizontal_margin, new_height), Image.ANTIALIAS)
resized_mask = distorted_mask.resize(
(size - horizontal_margin, new_height), Image.NEAREST)
background_width = size
background_height = new_height + vertical_margin
else:
raise ValueError("Invalid orientation")
#############################
# Generate background image #
#############################
if background_type == 0:
background_img = background_generator.gaussian_noise(
background_height, background_width)
elif background_type == 1:
background_img = background_generator.plain_white(
background_height, background_width)
elif background_type == 2:
background_img = background_generator.quasicrystal(
background_height, background_width)
else:
background_img = background_generator.image(
background_height, background_width, image_dir)
background_mask = Image.new("RGB",
(background_width, background_height),
(0, 0, 0))
#############################
# Place text with alignment #
#############################
new_text_width, _ = resized_img.size
if alignment == 0 or width == -1:
background_img.paste(resized_img, (margin_left, margin_top),
resized_img)
background_mask.paste(resized_mask, (margin_left, margin_top))
elif alignment == 1:
background_img.paste(
resized_img,
(int(background_width / 2 - new_text_width / 2), margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(int(background_width / 2 - new_text_width / 2), margin_top),
)
else:
background_img.paste(
resized_img,
(background_width - new_text_width - margin_right, margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(background_width - new_text_width - margin_right, margin_top),
)
##################################
# Apply gaussian blur #
##################################
gaussian_filter = ImageFilter.GaussianBlur(
radius=blur if not random_blur else rnd.randint(0, blur))
final_image = background_img.filter(gaussian_filter)
final_mask = background_mask.filter(gaussian_filter)
#####################################
# Generate name for resulting image #
#####################################
if name_format == 0:
image_name = "{}_{}.{}".format(text, str(index), extension)
mask_name = "{}_{}_mask.png".format(text, str(index))
elif name_format == 1:
image_name = "{}_{}.{}".format(str(index), text, extension)
mask_name = "{}_{}_mask.png".format(str(index), text)
elif name_format == 2:
image_name = "{}.{}".format(str(index), extension)
mask_name = "{}_mask.png".format(str(index))
else:
print("{} is not a valid name format. Using default.".format(
name_format))
image_name = "{}_{}.{}".format(text, str(index), extension)
mask_name = "{}_{}_mask.png".format(text, str(index))
# Save the image
if out_dir is not None:
final_image.convert("RGB").save(os.path.join(out_dir, image_name))
if output_mask == 1:
final_mask.convert("RGB").save(os.path.join(
out_dir, mask_name))
else:
if output_mask == 1:
return final_image.convert("RGB"), final_mask.convert("RGB")
return final_image.convert("RGB")
def generate(
cls,
index,
text,
font,
out_dir,
size,
extension,
skewing_angle,
random_skew,
blur,
random_blur,
background_type,
distorsion_type,
distorsion_orientation,
is_handwritten,
name_format,
width,
alignment,
text_color,
orientation,
space_width,
character_spacing=0,
margins=(5, 5, 5, 5),
fit=False,
output_mask=False,
word_split=False,
image_dir=None,
stroke_width=0,
stroke_fill="#282828",
image_mode="RGB",
):
image = None
margin_top, margin_left, margin_bottom, margin_right = margins
horizontal_margin = margin_left + margin_right
vertical_margin = margin_top + margin_bottom
##########################
# Create picture of text #
##########################
image, _, bboxes = generate(
text,
font,
text_color,
size,
orientation,
space_width,
character_spacing,
fit,
word_split,
stroke_width,
stroke_fill,
)
##################################
# Resize image to desired format #
##################################
# Horizontal text
if orientation == 0:
ratio = float(size - vertical_margin) / float(image.size[1])
bboxes = bbox_resize(bboxes, (ratio, ratio))
# ratio_y = size - vertical_margin / size
new_width = int(image.size[0] * ratio)
resized_img = image.resize((new_width, size - vertical_margin),
Image.ANTIALIAS)
background_width = width if width > 0 else new_width + horizontal_margin
background_height = size
# Vertical text
elif orientation == 1:
ratio = (float(size - horizontal_margin) / float(image.size[0]))
bboxes = bbox_resize(bboxes, (ratio, ratio))
new_height = int(float(image.size[1]) * ratio)
resized_img = image.resize((size - horizontal_margin, new_height),
Image.ANTIALIAS)
background_width = size
background_height = new_height + vertical_margin
else:
raise ValueError("Invalid orientation")
#############################
# Generate background image #
#############################
if background_type == 0:
background = background_generator.gaussian_noise(
background_height, background_width)
elif background_type == 1:
background = background_generator.plain_white(
background_height, background_width)
elif background_type == 2:
background = background_generator.quasicrystal(
background_height, background_width)
else:
background = background_generator.picture(background_height,
background_width)
#############################
# Place text with alignment #
#############################
new_text_width, _ = resized_img.size
if alignment == 0 or width == -1:
background.paste(resized_img, (margin_left, margin_top),
resized_img)
bboxes = bbox_margin(bboxes, (margin_left, margin_top))
elif alignment == 1:
background.paste(
resized_img,
(int(background_width / 2 - new_text_width / 2), margin_top),
resized_img,
)
else:
background.paste(
resized_img,
(background_width - new_text_width - margin_right, margin_top),
resized_img,
)
##################################
# Apply gaussian blur #
##################################
final_image = background.filter(
ImageFilter.GaussianBlur(
radius=(blur if not random_blur else rnd.randint(0, blur))))
#####################################
# Generate name for resulting image #
#####################################
if name_format == 0:
image_name = "{}_{}.{}".format(text, str(index), extension)
elif name_format == 1:
image_name = "{}_{}.{}".format(str(index), text, extension)
elif name_format == 2:
image_name = "{}.{}".format(str(index), extension)
else:
print("{} is not a valid name format. Using default.".format(
name_format))
image_name = "{}_{}.{}".format(text, str(index), extension)
# Save the image
if out_dir is not None:
final_image.convert("RGB").save(os.path.join(out_dir, image_name))
else:
return final_image.convert("RGB"), bboxes
def generate(
cls,
index,
text,
font,
out_dir,
size,
extension,
skewing_angle,
random_skew,
blur,
random_blur,
background_type,
distorsion_type,
distorsion_orientation,
is_handwritten,
name_format,
width,
alignment,
text_color,
orientation,
space_width,
character_spacing,
margins,
fit,
output_mask,
word_split,
image_dir,
stroke_width=0,
stroke_fill="#282828",
image_mode="RGB",
output_bboxes=0,
):
image = None
margin_top, margin_left, margin_bottom, margin_right = margins
horizontal_margin = margin_left + margin_right
vertical_margin = margin_top + margin_bottom
##########################
# Create picture of text #
##########################
if is_handwritten:
if orientation == 1:
raise ValueError("Vertical handwritten text is unavailable")
image, mask = handwritten_text_generator.generate(text, text_color)
else:
image, mask = computer_text_generator.generate(
text,
font,
text_color,
size,
orientation,
space_width,
character_spacing,
fit,
word_split,
stroke_width,
stroke_fill,
)
random_angle = rnd.randint(0 - skewing_angle, skewing_angle)
rotated_img = image.rotate(
skewing_angle if not random_skew else random_angle, expand=1)
rotated_mask = mask.rotate(
skewing_angle if not random_skew else random_angle, expand=1)
#############################
# Apply distorsion to image #
#############################
if distorsion_type == 0:
distorted_img = rotated_img # Mind = blown
distorted_mask = rotated_mask
elif distorsion_type == 1:
distorted_img, distorted_mask = distorsion_generator.sin(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
elif distorsion_type == 2:
distorted_img, distorted_mask = distorsion_generator.cos(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
else:
distorted_img, distorted_mask = distorsion_generator.random(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
##################################
# Resize image to desired format #
##################################
# Horizontal text
if orientation == 0:
new_width = int(
distorted_img.size[0] *
(float(size - vertical_margin) / float(distorted_img.size[1])))
resized_img = distorted_img.resize(
(new_width, size - vertical_margin), Image.ANTIALIAS)
resized_mask = distorted_mask.resize(
(new_width, size - vertical_margin), Image.NEAREST)
background_width = width if width > 0 else new_width + horizontal_margin
background_height = size
# Vertical text
elif orientation == 1:
new_height = int(
float(distorted_img.size[1]) *
(float(size - horizontal_margin) /
float(distorted_img.size[0])))
resized_img = distorted_img.resize(
(size - horizontal_margin, new_height), Image.ANTIALIAS)
resized_mask = distorted_mask.resize(
(size - horizontal_margin, new_height), Image.NEAREST)
background_width = size
background_height = new_height + vertical_margin
else:
raise ValueError("Invalid orientation")
#############################
# Generate background image #
#############################
if background_type == 0:
background_img = background_generator.gaussian_noise(
background_height, background_width)
elif background_type == 1:
background_img = background_generator.plain_white(
background_height, background_width)
elif background_type == 2:
background_img = background_generator.quasicrystal(
background_height, background_width)
else:
background_img = background_generator.image(
background_height, background_width, image_dir)
background_mask = Image.new("RGB",
(background_width, background_height),
(0, 0, 0))
##############################################################
# Comparing average pixel value of text and background image #
##############################################################
try:
resized_img_st = ImageStat.Stat(resized_img,
resized_mask.split()[2])
background_img_st = ImageStat.Stat(background_img)
resized_img_px_mean = sum(resized_img_st.mean[:2]) / 3
background_img_px_mean = sum(background_img_st.mean) / 3
if abs(resized_img_px_mean - background_img_px_mean) < 15:
print("value of mean pixel is too similar. Ignore this image")
print("resized_img_st \n {}".format(resized_img_st.mean))
print("background_img_st \n {}".format(background_img_st.mean))
return
except Exception as err:
return
#############################
# Place text with alignment #
#############################
new_text_width, _ = resized_img.size
if alignment == 0 or width == -1:
background_img.paste(resized_img, (margin_left, margin_top),
resized_img)
background_mask.paste(resized_mask, (margin_left, margin_top))
elif alignment == 1:
background_img.paste(
resized_img,
(int(background_width / 2 - new_text_width / 2), margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(int(background_width / 2 - new_text_width / 2), margin_top),
)
else:
background_img.paste(
resized_img,
(background_width - new_text_width - margin_right, margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(background_width - new_text_width - margin_right, margin_top),
)
#######################
# Apply gaussian blur #
#######################
gaussian_filter = ImageFilter.GaussianBlur(
radius=blur if not random_blur else rnd.randint(0, blur))
final_image = background_img.filter(gaussian_filter)
final_mask = background_mask.filter(gaussian_filter)
############################################
# Change image mode (RGB, grayscale, etc.) #
############################################
final_image = final_image.convert(image_mode)
final_mask = final_mask.convert(image_mode)
#####################################
# Generate name for resulting image #
#####################################
# We remove spaces if space_width == 0
if space_width == 0:
text = text.replace(" ", "")
if name_format == 0:
name = "{}_{}".format(text, str(index))
elif name_format == 1:
name = "{}_{}".format(str(index), text)
elif name_format == 2:
name = str(index)
else:
print("{} is not a valid name format. Using default.".format(
name_format))
name = "{}_{}".format(text, str(index))
image_name = "{}.{}".format(name, extension)
mask_name = "{}_mask.png".format(name)
box_name = "{}_boxes.txt".format(name)
tess_box_name = "{}.box".format(name)
# Save the image
if out_dir is not None:
final_image.save(os.path.join(out_dir, image_name))
if output_mask == 1:
final_mask.save(os.path.join(out_dir, mask_name))
if output_bboxes == 1:
bboxes = mask_to_bboxes(final_mask)
with open(os.path.join(out_dir, box_name), "w") as f:
for bbox in bboxes:
f.write(" ".join([str(v) for v in bbox]) + "\n")
if output_bboxes == 2:
bboxes = mask_to_bboxes(final_mask, tess=True)
with open(os.path.join(out_dir, tess_box_name), "w") as f:
for bbox, char in zip(bboxes, text):
f.write(" ".join([char] + [str(v)
for v in bbox] + ['0']) +
"\n")
else:
if output_mask == 1:
return final_image, final_mask
return final_image
def generate(cls,field,index,text,font,out_dir,size,extension,skewing_angle,random_skew,
blur,random_blur,background_type,distorsion_type,distorsion_orientation,
is_handwritten,name_format,width,alignment,text_color,orientation,space_width,
character_spacing,margins,fit,output_mask,word_split,image_dir,stroke_width=0,
stroke_fill="#282828",image_mode="RGB",
):
image = None
margin_top, margin_left, margin_bottom, margin_right = margins
horizontal_margin = margin_left + margin_right
vertical_margin = margin_top + margin_bottom
##########################
# Create picture of text #
##########################
if is_handwritten:
if orientation == 1:
raise ValueError("Vertical handwritten text is unavailable")
image, mask = handwritten_text_generator.generate(text, text_color)
else:
image, mask = computer_text_generator.generate(text,font,text_color,size,orientation,
space_width,character_spacing,fit,word_split,
stroke_width, stroke_fill,)
random_angle = rnd.randint(0 - skewing_angle, skewing_angle)
rotated_img = image.rotate(
skewing_angle if not random_skew else random_angle, expand=1
)
rotated_mask = mask.rotate(
skewing_angle if not random_skew else random_angle, expand=1
)
#############################
# Apply distorsion to image #
#############################
if distorsion_type == 0:
distorted_img = rotated_img # Mind = blown
distorted_mask = rotated_mask
elif distorsion_type == 1:
distorted_img, distorted_mask = distorsion_generator.sin(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0 or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1 or distorsion_orientation == 2),
)
elif distorsion_type == 2:
distorted_img, distorted_mask = distorsion_generator.cos(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0 or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1 or distorsion_orientation == 2),
)
else:
distorted_img, distorted_mask = distorsion_generator.random(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0 or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1 or distorsion_orientation == 2),
)
##################################
# Resize image to desired format #
##################################
# Horizontal text
if orientation == 0:
new_width = int(
distorted_img.size[0]
* (float(size - vertical_margin) / float(distorted_img.size[1]))
)
resized_img = distorted_img.resize(
(new_width, size - vertical_margin), Image.ANTIALIAS
)
resized_mask = distorted_mask.resize((new_width, size - vertical_margin), Image.NEAREST)
background_width = width if width > 0 else new_width + horizontal_margin
background_height = size
# Vertical text
elif orientation == 1:
new_height = int(
float(distorted_img.size[1])
* (float(size - horizontal_margin) / float(distorted_img.size[0]))
)
resized_img = distorted_img.resize(
(size - horizontal_margin, new_height), Image.ANTIALIAS
)
resized_mask = distorted_mask.resize(
(size - horizontal_margin, new_height), Image.NEAREST
)
background_width = size
background_height = new_height + vertical_margin
else:
raise ValueError("Invalid orientation")
#############################
# Generate background image #
#############################
if background_type == 0:
background_img = background_generator.gaussian_noise(background_height, background_width)
elif background_type == 1:
background_img = background_generator.plain_white(background_height, background_width)
elif background_type == 2:
background_img = background_generator.quasicrystal(background_height, background_width)
else:
background_img = background_generator.image(background_height, background_width, image_dir)
background_img = change_brightness(background_img, value=rnd.randrange(0,30,5))
background_mask = Image.new("RGB", (background_width, background_height), (0, 0, 0))
#############################
# Place text with alignment #
#############################
resized_img = add_noise(resized_img)
resized_img = resized_img.filter(ImageFilter.BLUR)
new_text_width, _ = resized_img.size
issue_loc = [(325, 50), (310, 42), (320, 58), (311, 45)]
if alignment == 0 or width == -1:
background_img.paste(resized_img, (margin_left, margin_top), resized_img)
background_mask.paste(resized_mask, (margin_left, margin_top))
elif alignment == 1:
list_img = os.listdir(image_dir)
l = len(list_img)
m = rnd.randint(0, l-1)
img_path = image_dir + "/" + list_img[m]
init_img = Image.open(img_path)
bg_width, bg_height = init_img.size
background_img = background_generator.image(bg_height-10, bg_width-5, image_dir)
background_img = change_brightness(background_img, value=rnd.randrange(10,30,5))
background_mask = Image.new("RGB", (bg_width, bg_height), (0, 0, 0))
k = rnd.randint(0, 3)
background_img.paste(resized_img,issue_loc[k],resized_img,)
background_mask.paste(resized_mask,(int(background_width / 2 - new_text_width / 2), margin_top),)
else:
background_img.paste(
resized_img,
(background_width - new_text_width - margin_right, margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(background_width - new_text_width - margin_right, margin_top),
)
#######################
# Apply gaussian blur #
#######################
# gaussian_filter = ImageFilter.GaussianBlur(
# radius=blur if not random_blur else rnd.randint(1, blur)
# )
# final_image = background_img.filter(gaussian_filter)
# final_mask = background_mask.filter(gaussian_filter)
############################################
# Change image mode (RGB, grayscale, etc.) #
############################################
final_image = background_img.convert(image_mode)
final_mask = background_mask.convert(image_mode)
#####################################
# Generate name for resulting image #
#####################################
# We remove spaces if space_width == 0
if space_width == 0:
text = text.replace(" ", "")
if name_format == 0:
image_name = "{}_{:05d}.{}".format(field, index, extension)
mask_name = "{}_{}_mask.png".format(text, str(index))
elif name_format == 1:
image_name = "{}_{:05d}_0.{}".format(field, index, extension)
mask_name = "{}_{}_mask.png".format(str(index), text)
elif name_format == 2:
image_name = "{}_{:05d}_1.{}".format(field, index, extension)
mask_name = "{}_{}_mask.png".format(str(index), text)
elif name_format == 3:
image_name = "{}.{}".format(str(index), extension)
mask_name = "{}_mask.png".format(str(index))
else:
print("{} is not a valid name format. Using default.".format(name_format))
image_name = "{}_{}.{}".format(text, str(index), extension)
mask_name = "{}_{}_mask.png".format(text, str(index))
# Save the image
if out_dir is not None:
final_image.save(os.path.join(out_dir, image_name))
if output_mask == 1:
final_mask.save(os.path.join(out_dir, mask_name))
else:
if output_mask == 1:
return final_image, final_mask
return final_image
def generate(
cls,
index,
strdict,
font,
out_dir,
size,
extension,
skewing_angle,
random_skew,
blur,
random_blur,
background_type,
distorsion_type,
distorsion_orientation,
is_handwritten,
name_format,
width,
alignment,
text_color,
orientation,
space_width,
character_spacing,
margins,
fit,
output_mask,
word_split,
image_dir,
stroke_width=0,
stroke_fill="#282828",
image_mode="RGB",
):
background_height = 800
background_width = 600
row_number = 25
row_number_range = [i + 1 for i in range(row_number)]
row_height = int(background_height / row_number)
row_text_maxlength = 40
font_size_range = list(range(int(row_height - 12),
int(row_height + 1)))
#############################
# Generate background image #
#############################
if background_type == 0:
background_img = background_generator.gaussian_noise(
background_height, background_width)
elif background_type == 1:
background_img = background_generator.plain_white(
background_height, background_width)
elif background_type == 2:
background_img = background_generator.quasicrystal(
background_height, background_width)
else:
background_img = background_generator.image(
background_height, background_width, image_dir)
# background_mask = Image.new(
# "RGB", (background_width, background_height), (0, 0, 0)
# )
bbox = []
texts = []
row_hastxt = []
for _ in range(rnd.choice(row_number_range)):
rowid = rnd.choice(row_number_range) - 1
while rowid in row_hastxt:
rowid = rnd.choice(row_number_range) - 1
row_hastxt.append(rowid)
##########################
# Create picture of text #
##########################
text = create_strings_from_dict(row_text_maxlength, strdict)
texts.append(text)
font_size = rnd.choice(font_size_range)
image, mask = computer_text_generator.generate(
text,
font,
text_color,
font_size,
orientation,
space_width,
character_spacing,
fit,
word_split,
stroke_width,
stroke_fill,
)
random_angle = rnd.randint(0 - skewing_angle, skewing_angle)
rotated_img = image.rotate(
skewing_angle if not random_skew else random_angle, expand=1)
rotated_mask = mask.rotate(
skewing_angle if not random_skew else random_angle, expand=1)
#############################
# Apply distorsion to image #
#############################
if distorsion_type == 0:
distorted_img = rotated_img # Mind = blown
# distorted_mask = rotated_mask
elif distorsion_type == 1:
distorted_img, distorted_mask = distorsion_generator.sin(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
elif distorsion_type == 2:
distorted_img, distorted_mask = distorsion_generator.cos(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
else:
distorted_img, distorted_mask = distorsion_generator.random(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
##################################
# Resize image to desired format #
##################################
# Horizontal text
if orientation == 0:
new_width = int(
distorted_img.size[0] *
(float(font_size) / float(distorted_img.size[1])))
if new_width <= background_width:
resized_img = distorted_img.resize((new_width, font_size),
Image.ANTIALIAS)
# resized_mask = distorted_mask.resize((new_width, font_size), Image.NEAREST)
else:
new_height = int(distorted_img.size[1] *
(float(background_width) /
float(distorted_img.size[0])))
resized_img = distorted_img.resize(
(background_width, new_height), Image.ANTIALIAS)
resized_mask = distorted_img.resize(
(background_width, new_height), Image.NEAREST)
# Vertical text
elif orientation == 1:
# new_height = int(
# float(distorted_img.size[1])
# * (float(size - horizontal_margin) / float(distorted_img.size[0]))
# )
# resized_img = distorted_img.resize(
# (size - horizontal_margin, new_height), Image.ANTIALIAS
# )
# resized_mask = distorted_mask.resize(
# (size - horizontal_margin, new_height), Image.NEAREST
# )
# background_width = size
# background_height = new_height + vertical_margin
raise ValueError("Vertical text")
else:
raise ValueError("Invalid orientation")
#############################
# Place text with alignment #
#############################
top_left_x = rnd.choice(
range(0, max(1, background_width - resized_img.size[0])))
top_left_y = (rowid * row_height) + rnd.choice(
range(0, max(1, row_height - font_size)))
bbox.append([
top_left_x, top_left_y, top_left_x + resized_img.size[0] - 1,
top_left_y + resized_img.size[1] - 1
])
background_img.paste(
resized_img,
(top_left_x, top_left_y),
resized_img,
)
#######################
# Apply gaussian blur #
#######################
gaussian_filter = ImageFilter.GaussianBlur(
radius=blur if not random_blur else rnd.randint(0, blur))
final_image = background_img.filter(gaussian_filter)
# final_mask = background_mask.filter(gaussian_filter)
############################################
# Change image mode (RGB, grayscale, etc.) #
############################################
final_image = final_image.convert(image_mode)
# final_mask = final_mask.convert(image_mode)
#####################################
# Generate name for resulting image #
#####################################
# We remove spaces if space_width == 0
if space_width == 0:
text = text.replace(" ", "")
if name_format == 0:
image_name = "{}_{}.{}".format(text, str(index), extension)
mask_name = "{}_{}_mask.png".format(text, str(index))
elif name_format == 1:
image_name = "{}_{}.{}".format(str(index), text, extension)
mask_name = "{}_{}_mask.png".format(str(index), text)
elif name_format == 2:
image_name = "{}.{}".format(str(index), extension)
mask_name = "{}_mask.png".format(str(index))
else:
print("{} is not a valid name format. Using default.".format(
name_format))
image_name = "{}_{}.{}".format(text, str(index), extension)
mask_name = "{}_{}_mask.png".format(text, str(index))
# Save the image
if out_dir is not None:
final_image.save(os.path.join(out_dir, image_name))
# if output_mask == 1:
# final_mask.save(os.path.join(out_dir, mask_name))
else:
# if output_mask == 1:
# return final_image, final_mask
return final_image, bbox, texts
def generate_img(
index,
text,
font,
out_dir,
size,
extension,
skewing_angle,
random_skew,
blur,
random_blur,
background_type,
distorsion_type,
distorsion_orientation,
is_handwritten,
name_format,
width,
alignment,
text_color,
orientation,
space_width,
character_spacing,
margins,
fit,
output_mask,
word_split,
image_dir,
stroke_width=0,
stroke_fill="#282828",
image_mode="RGB",
):
image = None
margin_top, margin_left, margin_bottom, margin_right = margins
horizontal_margin = margin_left + margin_right
vertical_margin = margin_top + margin_bottom
##########################
# Create picture of text #
##########################
image, mask, bboxes = generate(
text,
font,
text_color,
size,
orientation,
space_width,
character_spacing,
fit,
word_split,
stroke_width,
stroke_fill,
)
#############################
# Apply distorsion to image #
#############################
if distorsion_type == 0:
distorted_img = image # Mind = blown
distorted_mask = mask
else:
raise('Not use distorsion')
##################################
# Resize image to desired format #
##################################
# Horizontal text
if orientation == 0:
ratio = float(size - vertical_margin) / float(image.size[1])
bboxes = bbox_resize(bboxes, (ratio, ratio))
new_width = int(
distorted_img.size[0]
* (float(size - vertical_margin) / float(distorted_img.size[1]))
)
resized_img = distorted_img.resize(
(new_width, size - vertical_margin), Image.ANTIALIAS
)
resized_mask = distorted_mask.resize((new_width, size - vertical_margin), Image.NEAREST)
background_width = width if width > 0 else new_width + horizontal_margin
background_height = size
# Vertical text
elif orientation == 1:
ratio = (float(size - horizontal_margin) / float(image.size[0]))
bboxes = bbox_resize(bboxes, (ratio, ratio))
new_height = int(
float(distorted_img.size[1])
* (float(size - horizontal_margin) / float(distorted_img.size[0]))
)
resized_img = distorted_img.resize(
(size - horizontal_margin, new_height), Image.ANTIALIAS
)
resized_mask = distorted_mask.resize(
(size - horizontal_margin, new_height), Image.NEAREST
)
background_width = size
background_height = new_height + vertical_margin
else:
raise ValueError("Invalid orientation")
#############################
# Generate background image #
#############################
if background_type == 0:
background_img = background_generator.gaussian_noise(
background_height, background_width
)
elif background_type == 1:
background_img = background_generator.plain_white(
background_height, background_width
)
elif background_type == 2:
background_img = Image.new(
"RGB", (background_width, background_height), stroke_fill
)
else:
background_img = background_generator.image(
background_height, background_width, image_dir
)
background_mask = Image.new(
"RGB", (background_width, background_height), (0, 0, 0)
)
#############################
# Place text with alignment #
#############################
new_text_width, _ = resized_img.size
if alignment == 0 or width == -1:
background_img.paste(resized_img, (margin_left, margin_top), resized_img)
background_mask.paste(resized_mask, (margin_left, margin_top))
bboxes = bbox_margin(bboxes, (margin_left, margin_top))
elif alignment == 1:
background_img.paste(
resized_img,
(int(background_width / 2 - new_text_width / 2), margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(int(background_width / 2 - new_text_width / 2), margin_top),
)
else:
background_img.paste(
resized_img,
(background_width - new_text_width - margin_right, margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(background_width - new_text_width - margin_right, margin_top),
)
#######################
# Apply gaussian blur #
#######################
gaussian_filter = ImageFilter.GaussianBlur(
radius=blur if not random_blur else rnd.randint(0, blur)
)
final_image = background_img.filter(gaussian_filter)
final_mask = background_mask.filter(gaussian_filter)
############################################
# Change image mode (RGB, grayscale, etc.) #
############################################
final_image = final_image.convert(image_mode)
final_mask = final_mask.convert(image_mode)
if output_mask == 1:
return final_image, final_mask
return final_image, bboxes
def generate(
cls,
index,
text,
font,
out_dir,
size,
extension,
skewing_angle,
random_skew,
blur,
random_blur,
background_type,
distorsion_type,
distorsion_orientation,
is_handwritten,
name_format,
width,
alignment,
text_color,
orientation,
space_width,
character_spacing,
margins,
fit,
output_mask,
word_split,
image_dir,
bounding_box,
draw_bbox,
radius,
):
image = None
margin_top, margin_left, margin_bottom, margin_right = margins
horizontal_margin = margin_left + margin_right
vertical_margin = margin_top + margin_bottom
##########################
# Create picture of text #
##########################
if is_handwritten:
if orientation == 1:
raise ValueError("Vertical handwritten text is unavailable")
image, mask = handwritten_text_generator.generate(text, text_color)
else:
image, mask, boxs = computer_text_generator.generate(
text,
font,
text_color,
size,
orientation,
space_width,
character_spacing,
fit,
word_split,
radius,
)
random_angle = rnd.randint(0 - skewing_angle, skewing_angle)
rotated_img = image.rotate(
skewing_angle if not random_skew else random_angle, expand=True)
rotated_mask = mask.rotate(
skewing_angle if not random_skew else random_angle, expand=True)
########## Written by xy.huang ##########
########## Rotate bounding_boxs##########
if bounding_box == True:
expand_matrix = np.stack([(rotated_img.width - image.width) / 2,
(rotated_img.height - image.height) / 2],
axis=0)
scaler = np.stack([rotated_img.width, rotated_img.height], axis=0)
center = np.reshape(0.5 * scaler, [1, 2])
angle = skewing_angle if not random_skew else random_angle
theta = (360 - angle) * (math.pi / 180.0)
rotation = np.stack(
[np.cos(theta),
np.sin(theta), -np.sin(theta),
np.cos(theta)],
axis=0)
rotation_matrix = np.reshape(rotation, [2, 2])
for i in range(len(boxs)):
box = np.reshape(boxs[i], (4, 2))
box = box + expand_matrix
box = np.matmul(box - center, rotation_matrix) + center
boxs[i] = np.reshape(box, (8, )).tolist()
########## Written by xy.huang ##########
#############################
# Apply distorsion to image #
#############################
if distorsion_type == 0:
distorted_img = rotated_img # Mind = blown
distorted_mask = rotated_mask
elif distorsion_type == 1:
distorted_img, distorted_mask = distorsion_generator.sin(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
elif distorsion_type == 2:
distorted_img, distorted_mask = distorsion_generator.cos(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
else:
distorted_img, distorted_mask = distorsion_generator.random(
rotated_img,
rotated_mask,
vertical=(distorsion_orientation == 0
or distorsion_orientation == 2),
horizontal=(distorsion_orientation == 1
or distorsion_orientation == 2),
)
##################################
# Resize image to desired format #
##################################
original_width, original_height = distorted_img.size
# Horizontal text
if orientation == 0:
new_width = int(
distorted_img.size[0] *
(float(size - vertical_margin) / float(distorted_img.size[1])))
resized_img = distorted_img.resize(
(new_width, size - vertical_margin), Image.ANTIALIAS)
resized_mask = distorted_mask.resize(
(new_width, size - vertical_margin), Image.NEAREST)
background_width = width if width > 0 else new_width + horizontal_margin
background_height = size
if bounding_box == True:
width_scale = new_width / original_width
height_scale = (size - vertical_margin) / original_height
# Vertical text
elif orientation == 1:
new_height = int(
float(distorted_img.size[1]) *
(float(size - horizontal_margin) /
float(distorted_img.size[0])))
resized_img = distorted_img.resize(
(size - horizontal_margin, new_height), Image.ANTIALIAS)
resized_mask = distorted_mask.resize(
(size - horizontal_margin, new_height), Image.NEAREST)
background_width = size
background_height = new_height + vertical_margin
if bounding_box == True:
width_scale = (size - horizontal_margin) / original_width
height_scale = new_height / original_height
#Curved text
elif orientation == 2:
resized_img = distorted_img
resized_mask = distorted_img
background_width = resized_img.width
background_height = resized_img.height
if bounding_box == True:
width_scale = 1
height_scale = 1
else:
raise ValueError("Invalid orientation")
########## Written by xy.huang ##########
if bounding_box == True:
for box in boxs:
for i in range(8):
if i % 2 == 0:
box[i] = box[i] * width_scale
else:
box[i] = box[i] * height_scale
########## Written by xy.huang ##########
#############################
# Generate background image #
#############################
if background_type == 0:
background_img = background_generator.gaussian_noise(
background_height, background_width)
elif background_type == 1:
background_img = background_generator.plain_white(
background_height, background_width)
elif background_type == 2:
background_img = background_generator.quasicrystal(
background_height, background_width)
else:
background_img = background_generator.image(
background_height, background_width, image_dir)
background_mask = Image.new("RGB",
(background_width, background_height),
(0, 0, 0))
#############################
# Place text with alignment #
#############################
new_text_width, _ = resized_img.size
if alignment == 0 or width == -1:
background_img.paste(resized_img, (margin_left, margin_top),
resized_img)
background_mask.paste(resized_mask, (margin_left, margin_top))
background_offset_width = margin_left
elif alignment == 1:
background_img.paste(
resized_img,
(int(background_width / 2 - new_text_width / 2), margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(int(background_width / 2 - new_text_width / 2), margin_top),
)
background_offset_width = int(background_width / 2 -
new_text_width / 2)
else:
background_img.paste(
resized_img,
(background_width - new_text_width - margin_right, margin_top),
resized_img,
)
background_mask.paste(
resized_mask,
(background_width - new_text_width - margin_right, margin_top),
)
background_offset_width = background_width - new_text_width - margin_right
########## Written by xy.huang ##########
if bounding_box == True:
for box in boxs:
for i in range(8):
if i % 2 == 0:
box[i] = box[i] + background_offset_width
else:
box[i] = box[i] + margin_top
########## Written by xy.huang ##########
##################################
# Apply gaussian blur #
##################################
gaussian_filter = ImageFilter.GaussianBlur(
radius=blur if not random_blur else rnd.randint(0, blur))
final_image = background_img.filter(gaussian_filter)
final_mask = background_mask.filter(gaussian_filter)
#####################################
# Generate name for resulting image #
#####################################
if name_format == 0:
image_name = "{}_{}.{}".format(text, str(index), extension)
mask_name = "{}_{}_mask.png".format(text, str(index))
txt_name = "{}_{}.txt".format(text, str(index))
elif name_format == 1:
image_name = "{}_{}.{}".format(str(index), text, extension)
mask_name = "{}_{}_mask.png".format(str(index), text)
txt_name = "{}_{}.txt".format(str(index), text)
elif name_format == 2:
image_name = "{}.{}".format(str(index), extension)
mask_name = "{}_mask.png".format(str(index))
txt_name = "{}.txt".format(str(index))
else:
print("{} is not a valid name format. Using default.".format(
name_format))
image_name = "{}_{}.{}".format(text, str(index), extension)
mask_name = "{}_{}_mask.png".format(text, str(index))
txt_name = "{}_{}.txt".format(text, str(index))
# Draw coordinate
if draw_bbox == True:
if bounding_box != True:
raise ValueError(
"Draw bounding_box only works with -bb or --bounding_box")
bboxDraw = ImageDraw.Draw(final_image)
for box in boxs:
bboxDraw.polygon(box, outline=(255, 0, 0))
splitted_text = []
for p in text:
if p != ' ':
splitted_text.append(p)
# Save the image
if out_dir is not None:
final_image.convert("RGB").save(os.path.join(out_dir, image_name))
if output_mask == 1:
final_mask.convert("RGB").save(os.path.join(
out_dir, mask_name))
if bounding_box == True:
with open(os.path.join(out_dir, txt_name), mode='w') as f:
for box, p in zip(boxs, splitted_text):
f.writelines(str(box) + p)
f.write('\r\n')
else:
if output_mask == 1:
#return final_image.convert("RGBA"), final_mask.convert("RGB"), boxs
return final_image, final_mask, boxs
#return final_image.convert("RGBA"), boxs
return final_image, boxs
