Spaces:
Running on Zero
Running on Zero
File size: 5,429 Bytes
b411c37 98e1d9f b411c37 98e1d9f b411c37 98e1d9f b411c37 98e1d9f b411c37 98e1d9f b411c37 98e1d9f b411c37 98e1d9f b411c37 98e1d9f b411c37 98e1d9f b411c37 | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | from __future__ import annotations
from io import BytesIO
from pathlib import Path
from PIL import Image
from reportlab.lib.utils import ImageReader
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfgen.canvas import Canvas
from ocr_studio.config import ARABIC_FONT_PATH, LATIN_FONT_PATH, PDF_DPI
from ocr_studio.language import contains_arabic_script
from ocr_studio.spotting import TextSpan
_FONTS_REGISTERED = False
LATIN_FONT = "NotoSans"
ARABIC_FONT = "NotoNaskhArabic"
def _register_fonts() -> None:
global _FONTS_REGISTERED
if _FONTS_REGISTERED:
return
if LATIN_FONT_PATH.exists():
pdfmetrics.registerFont(TTFont(LATIN_FONT, str(LATIN_FONT_PATH)))
if ARABIC_FONT_PATH.exists():
pdfmetrics.registerFont(TTFont(ARABIC_FONT, str(ARABIC_FONT_PATH)))
_FONTS_REGISTERED = True
def _font_for(text: str) -> str:
if contains_arabic_script(text) and ARABIC_FONT_PATH.exists():
return ARABIC_FONT
if LATIN_FONT_PATH.exists():
return LATIN_FONT
return "Helvetica"
def _fallback_spans(image: Image.Image, text: str, rtl: bool) -> list[TextSpan]:
lines = [line.strip() for line in (text or "").splitlines() if line.strip()]
if not lines:
return []
top = image.height * 0.06
bottom = image.height * 0.94
usable = max(bottom - top, float(len(lines) * 14))
line_height = usable / max(len(lines), 1)
left = image.width * 0.05
right = image.width * 0.95
spans: list[TextSpan] = []
for index, line in enumerate(lines):
y0 = top + index * line_height
y1 = min(image.height - 2.0, y0 + max(12.0, line_height * 0.85))
if rtl and contains_arabic_script(line):
box = (left + (right - left) * 0.15, y0, right, y1)
else:
box = (left, y0, right, y1)
spans.append(TextSpan(text=line, box=box))
return spans
def _draw_page(
canvas: Canvas,
image: Image.Image,
text: str,
spans: list[TextSpan] | None,
rtl: bool,
) -> None:
page_w = image.width * 72.0 / PDF_DPI
page_h = image.height * 72.0 / PDF_DPI
scale = page_w / float(image.width)
canvas.setPageSize((page_w, page_h))
buffer = BytesIO()
image.convert("RGB").save(buffer, format="JPEG", quality=92)
buffer.seek(0)
canvas.drawImage(
ImageReader(buffer),
0,
0,
width=page_w,
height=page_h,
preserveAspectRatio=True,
mask="auto",
)
overlay = [span for span in (spans or []) if span.text]
if not overlay:
overlay = _fallback_spans(image, text, rtl)
for span in overlay:
box = span.box or (image.width * 0.05, image.height * 0.05, image.width * 0.95, image.height * 0.12)
x0, y0, x1, y1 = box
pdf_x = x0 * scale
pdf_y = page_h - (y1 * scale)
box_h = max(8.0, (y1 - y0) * scale)
box_w = max(12.0, (x1 - x0) * scale)
font_size = max(6.0, min(box_h * 0.82, 28.0))
font_name = _font_for(span.text)
canvas.setFont(font_name, font_size)
canvas.setFillGray(0)
text_obj = canvas.beginText()
text_obj.setTextRenderMode(3)
text_obj.setTextOrigin(pdf_x, pdf_y + max(1.0, (box_h - font_size) / 2.0))
clipped = span.text
try:
text_width = canvas.stringWidth(clipped, font_name, font_size)
if text_width > box_w and box_w > 0:
font_size = max(5.0, font_size * (box_w / text_width))
canvas.setFont(font_name, font_size)
text_obj = canvas.beginText()
text_obj.setTextRenderMode(3)
text_obj.setTextOrigin(pdf_x, pdf_y + max(1.0, (box_h - font_size) / 2.0))
except Exception:
pass
text_obj.textOut(clipped)
canvas.drawText(text_obj)
def build_searchable_pdf(
image: Image.Image | list[Image.Image],
text: str | list[str],
destination: Path,
spans: list[TextSpan] | list[list[TextSpan]] | None = None,
rtl: bool = False,
) -> Path:
_register_fonts()
images = image if isinstance(image, list) else [image]
texts = text if isinstance(text, list) else [text]
if len(texts) < len(images):
texts = texts + [""] * (len(images) - len(texts))
span_pages: list[list[TextSpan]]
if spans is None:
span_pages = [[] for _ in images]
elif spans and isinstance(spans[0], TextSpan):
span_pages = [spans] # type: ignore[list-item]
if len(images) > 1:
span_pages = span_pages + [[] for _ in range(len(images) - 1)]
else:
span_pages = list(spans) # type: ignore[arg-type]
while len(span_pages) < len(images):
span_pages.append([])
first = images[0]
page_w = first.width * 72.0 / PDF_DPI
page_h = first.height * 72.0 / PDF_DPI
canvas = Canvas(str(destination), pagesize=(page_w, page_h))
canvas.setTitle("Searchable OCR scan")
canvas.setAuthor("Alireza Aminzadeh")
for index, page_image in enumerate(images):
if index:
canvas.showPage()
_draw_page(canvas, page_image, texts[index], span_pages[index], rtl)
canvas.save()
return destination
|