File size: 4,382 Bytes
2720487 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
import os
from typing import List, Tuple
import requests
from PIL import Image, ImageDraw, ImageFont
from surya.postprocessing.fonts import get_font_path
from surya.schema import TextLine
from surya.settings import settings
from surya.postprocessing.math.latex import is_latex
def sort_text_lines(lines: List[TextLine], tolerance=1.25):
# Sorts in reading order. Not 100% accurate, this should only
# be used as a starting point for more advanced sorting.
vertical_groups = {}
for line in lines:
group_key = round(line.bbox[1] / tolerance) * tolerance
if group_key not in vertical_groups:
vertical_groups[group_key] = []
vertical_groups[group_key].append(line)
# Sort each group horizontally and flatten the groups into a single list
sorted_lines = []
for _, group in sorted(vertical_groups.items()):
sorted_group = sorted(group, key=lambda x: x.bbox[0])
sorted_lines.extend(sorted_group)
return sorted_lines
def truncate_repetitions(text: str, min_len=15):
# From nougat, with some cleanup
if len(text) < 2 * min_len:
return text
# try to find a length at which the tail is repeating
max_rep_len = None
for rep_len in range(min_len, int(len(text) / 2)):
# check if there is a repetition at the end
same = True
for i in range(0, rep_len):
if text[len(text) - rep_len - i - 1] != text[len(text) - i - 1]:
same = False
break
if same:
max_rep_len = rep_len
if max_rep_len is None:
return text
lcs = text[-max_rep_len:]
# remove all but the last repetition
text_to_truncate = text
while text_to_truncate.endswith(lcs):
text_to_truncate = text_to_truncate[:-max_rep_len]
return text[:len(text_to_truncate)]
def get_text_size(text, font):
im = Image.new(mode="P", size=(0, 0))
draw = ImageDraw.Draw(im)
_, _, width, height = draw.textbbox((0, 0), text=text, font=font)
return width, height
def render_text(draw, text, s_bbox, bbox_width, bbox_height, font_path, box_font_size):
font = ImageFont.truetype(font_path, box_font_size)
text_width, text_height = get_text_size(text, font)
while (text_width > bbox_width or text_height > bbox_height) and box_font_size > 6:
box_font_size = box_font_size - 1
font = ImageFont.truetype(font_path, box_font_size)
text_width, text_height = get_text_size(text, font)
# Calculate text position (centered in bbox)
text_width, text_height = get_text_size(text, font)
x = s_bbox[0]
y = s_bbox[1] + (bbox_height - text_height) / 2
draw.text((x, y), text, fill="black", font=font)
def render_math(image, draw, text, s_bbox, bbox_width, bbox_height, font_path):
try:
from surya.postprocessing.math.render import latex_to_pil
box_font_size = max(10, min(int(.2 * bbox_height), 24))
img = latex_to_pil(text, bbox_width, bbox_height, fontsize=box_font_size)
img.thumbnail((bbox_width, bbox_height))
image.paste(img, (s_bbox[0], s_bbox[1]))
except Exception as e:
print(f"Failed to render math: {e}")
box_font_size = max(10, min(int(.75 * bbox_height), 24))
render_text(draw, text, s_bbox, bbox_width, bbox_height, font_path, box_font_size)
def draw_text_on_image(bboxes, texts, image_size: Tuple[int, int], langs: List[str], font_path=None, max_font_size=60, res_upscale=2, has_math=False):
if font_path is None:
font_path = get_font_path(langs)
new_image_size = (image_size[0] * res_upscale, image_size[1] * res_upscale)
image = Image.new('RGB', new_image_size, color='white')
draw = ImageDraw.Draw(image)
for bbox, text in zip(bboxes, texts):
s_bbox = [int(coord * res_upscale) for coord in bbox]
bbox_width = s_bbox[2] - s_bbox[0]
bbox_height = s_bbox[3] - s_bbox[1]
# Shrink the text to fit in the bbox if needed
if has_math and is_latex(text):
render_math(image, draw, text, s_bbox, bbox_width, bbox_height, font_path)
else:
box_font_size = max(6, min(int(.75 * bbox_height), max_font_size))
render_text(draw, text, s_bbox, bbox_width, bbox_height, font_path, box_font_size)
return image
|