cv-analysis-service/synthesis/segment/text_block.py

from typing import List

from PIL import Image, ImageDraw, ImageFont
from funcy import first

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.text_block_generator.paragraph import ParagraphGenerator
from synthesis.text.font import pick_random_mono_space_font_available_on_system


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


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