render_texts.py

import cv2
import json
import math
import random
from io import BytesIO
import numpy as np
import os
import concurrent.futures
import argparse
import pandas as pd
from PIL import ImageDraw, ImageFont, Image
from multiprocessing.pool import ThreadPool
from dataclasses import dataclass

@dataclass
class TextInstance:
    _font_mapping = None

    text: str
    width: float
    height: float
    left: float
    top: float
    angle: float
    font: str
    color: str = None

    # @property
    # def font_mapping(self):
    #     if type(self)._font_mapping is not None:
    #         font_files = pd.read_parquet("fonts.parquet")
    #         type(self)._font_mapping = font_files
    #     return type(self)._font_mapping

    def as_tuple(self):
        return (self.text, (self.left, self.top, self.width, self.height, self.angle))
    
    # def get_font_file(self):
    #     font_id, font_index = self.font.split(',')
    #     return self.font_mapping[(self.font_mapping["_id"] == font_id) & (self.font_mapping["index"] == int(font_index))]

    def get_font_file(self):
        font_id, font_index = self.font.split(',')
        return f"./fonts/{font_id}-{font_index}.ttf"

def render_text(image: Image, text_instance: TextInstance):
    text, bbox = text_instance.as_tuple()
    bbox = list(map(lambda x: x * 224 / 1645, bbox))
    font_file = text_instance.get_font_file()

    # Binary search for the largest font size that fits the bbox
    font_size = 1
    font_size_upper_bound = 200
    font_size_lower_bound = 1
    draw = ImageDraw.Draw(image)
    while font_size_lower_bound < font_size_upper_bound:
        font = ImageFont.truetype(font_file, size=font_size)
        left, top, right, bottom = draw.multiline_textbbox(bbox[:2], text, font=font, align="center")
        
        width = right - left
        height = bottom - top

        if width > bbox[2] or height > bbox[3]:
            font_size_upper_bound = font_size
        else:
            font_size_lower_bound = font_size + 1
        font_size = (font_size_lower_bound + font_size_upper_bound) // 2
    
    draw.multiline_text(
        bbox[:2],
        text,
        font=font,
        fill="#000000" if text_instance.color is None else text_instance.color,
        align="center",
    )

def main():
    parser = argparse.ArgumentParser()
    args = parser.parse_args()
    lines = []
    with os.scandir("1206_mistake_samples") as it:
        for entry in it:
            if entry.is_file() and entry.name.endswith(".json"):
                with open(entry.path, mode='r') as f:
                    meta = json.load(f)
                    img = Image.new("RGB", (224, 224), (0, 0, 0))
                    for text, style, bbox in zip(meta['texts'], meta['styles'], meta['bbox']):
                        text_instance = TextInstance(
                            text = text,
                            left = bbox[0],
                            top = bbox[1],
                            width = bbox[2],
                            height = bbox[3],
                            angle = bbox[4],
                            font = style["font-family"],
                            color = style.get("color", None)
                        )
                        render_text(img, text_instance)
                        img.save(f"./1206_mistake_samples/{entry.name}.png")

if __name__ == "__main__":
    main()