Matthias Bisping e258df899f Refactoring: Move
Move text block into its own module
2023-02-01 17:24:56 +01:00

113 lines
4.5 KiB
Python

import abc
import textwrap
from typing import List
from PIL import Image, ImageDraw, ImageFont
from faker import Faker
from funcy import first, identity, iterate, take, last, rest
from cv_analysis.utils import star, conj
from cv_analysis.utils.image_operations import superimpose
from cv_analysis.utils.rectangle import Rectangle
from synthesis.random import rnd
from synthesis.segment.content_rectangle import ContentRectangle
from synthesis.text.font import pick_random_mono_space_font_available_on_system
from synthesis.text.line_formatter.identity import IdentityLineFormatter
from synthesis.text.line_formatter.paragraph import ParagraphLineFormatter
class TextBlock(ContentRectangle):
def __init__(self, x1, y1, x2, y2, text_generator=None, font=None, font_size=None):
super().__init__(x1, y1, x2, y2)
self.font = font or ImageFont.load_default() # pick_random_font_available_on_system(size=font_size)
self.text_generator = text_generator or ParagraphGenerator()
def __call__(self, *args, **kwargs):
pass
def generate_random_text(self, rectangle: Rectangle, n_sentences=3000):
lines = self.text_generator(rectangle, n_sentences)
image = write_lines_to_image(lines, rectangle, self.font)
return self.__put_content(image)
def put_text(self, text: str, rectangle: Rectangle):
text_width, text_height = self.font.getsize(text)
width_delta = text_width - rectangle.width
height_delta = text_height - rectangle.height
image = Image.new("RGBA", (text_width, text_height), (0, 255, 255, 0))
if width_delta > 0 or height_delta > 0:
image = image.resize((int(rectangle.width * 0.9), text_height))
draw = ImageDraw.Draw(image)
draw.text((0, 0), text, font=self.font, fill=(0, 0, 0, 255))
return self.__put_content(image)
def __put_content(self, image: Image.Image):
self.content = image if not self.content else superimpose(self.content, image)
assert self.content.mode == "RGBA"
return self
class TextBlockGenerator(abc.ABC):
pass
class ParagraphGenerator(TextBlockGenerator):
def __init__(self):
self.line_formatter = ParagraphLineFormatter(blank_line_percentage=rnd.uniform(0, 0.5))
def __call__(self, rectangle, n_sentences):
return self.generate_paragraph(rectangle, n_sentences)
def generate_paragraph(self, rectangle, n_sentences):
lines = generate_random_text_lines(rectangle, self.line_formatter, n_sentences)
return lines
def write_lines_to_image(lines: List[str], rectangle: Rectangle, font=None) -> Image.Image:
def write_line(line, line_number):
draw.text((0, line_number * text_size), line, font=font, fill=(0, 0, 0, 255))
font = font or pick_random_mono_space_font_available_on_system()
image = Image.new("RGBA", (rectangle.width, rectangle.height), (0, 255, 255, 0))
draw = ImageDraw.Draw(image)
text_size = draw.textsize(first(lines), font=font)[1]
for line_number, line in enumerate(lines):
write_line(line, line_number)
return image
def generate_random_text_lines(rectangle: Rectangle, line_formatter=identity, n_sentences=3000) -> List[str]:
text = Faker().paragraph(nb_sentences=n_sentences, variable_nb_sentences=False, ext_word_list=None)
unformatted_lines = textwrap.wrap(text, width=rectangle.width, break_long_words=False)
# each iteration of the line formatter function formats one more line and adds it to the back of the list
formatted_lines_generator = iterate(star(line_formatter), (unformatted_lines, True))
# hence do as many iterations as there are lines in the rectangle
lines_per_iteration = take(len(unformatted_lines), formatted_lines_generator)
# and then take the lines from the last iteration of the function
formatted_lines, _ = last(lines_per_iteration)
return formatted_lines
class CaptionGenerator(TextBlockGenerator):
def __init__(self, caption_start=None):
self.line_formatter = IdentityLineFormatter()
self.caption_start = caption_start or f"Fig {rnd.randint(1, 20)}"
def __call__(self, rectangle, n_sentences):
return self.generate_paragraph(rectangle, n_sentences)
def generate_paragraph(self, rectangle, n_sentences):
lines = generate_random_text_lines(rectangle, self.line_formatter, n_sentences)
first_line_modified = f"{self.caption_start}.: {first(lines)}"
lines = conj(first_line_modified, rest(lines))
return lines