223d3e6ed0
- Add visitor support to content rectangle - Refactor box unpacking to happen by applying visitor pattern
78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
from math import ceil
|
|
from typing import List
|
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from funcy import first, take
|
|
|
|
from cv_analysis.logging import dev_logger, debug_log
|
|
from cv_analysis.utils.image_operations import superimpose
|
|
from cv_analysis.utils.rectangle import Rectangle
|
|
from synthesis.segment.content_rectangle import ContentRectangle
|
|
from synthesis.text.font import pick_random_mono_space_font_available_on_system
|
|
from synthesis.text.text_block_generator.paragraph import ParagraphGenerator
|
|
|
|
|
|
class TextBlock(ContentRectangle):
|
|
def __init__(self, x1, y1, x2, y2, text_generator=None, font=None, font_size=None):
|
|
dev_logger.trace(f"Creating text block at {x1, y1, x2, y2}.")
|
|
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()
|
|
|
|
@debug_log()
|
|
def __call__(self, *args, **kwargs):
|
|
pass
|
|
|
|
@debug_log()
|
|
def generate_random_text(self, rectangle: Rectangle, n_sentences=None):
|
|
lines = self.text_generator(rectangle, n_sentences)
|
|
image = write_lines_to_image(lines, rectangle, self.font)
|
|
return self.__put_content(image)
|
|
|
|
@debug_log()
|
|
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)
|
|
|
|
@debug_log()
|
|
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
|
|
|
|
|
|
@debug_log()
|
|
def write_lines_to_image(lines: List[str], rectangle: Rectangle, font=None) -> Image.Image:
|
|
@debug_log()
|
|
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]
|
|
|
|
# FIXME: This will sometimes cut off the last line, but it's better than underflowing, since then we would need to
|
|
# crop the target box to fit the text, and to the detection algorithm a cut off line makes little difference
|
|
# anyway. At least for now. If it should become relevant, we need to find a clean way to fit target box and text
|
|
# precisely.
|
|
n_lines = min(len(lines), ceil(rectangle.height / text_size))
|
|
|
|
for line_number, line in enumerate(take(n_lines, lines)):
|
|
write_line(line, line_number)
|
|
|
|
return image
|