From a5e57d381bc9630dbe39bf34b4b732e78e1282c0 Mon Sep 17 00:00:00 2001 From: felix Date: Thu, 30 Jul 2026 16:07:59 +0200 Subject: [PATCH 1/3] [reconstitution] Improve font scaling and rendering --- doctr/utils/reconstitution.py | 545 +++++++++++++++++++++++----------- 1 file changed, 370 insertions(+), 175 deletions(-) diff --git a/doctr/utils/reconstitution.py b/doctr/utils/reconstitution.py index b5afcc6210..acc734e976 100644 --- a/doctr/utils/reconstitution.py +++ b/doctr/utils/reconstitution.py @@ -5,7 +5,7 @@ import logging import math from functools import lru_cache -from typing import Any +from typing import Any, NamedTuple import numpy as np from anyascii import anyascii @@ -16,10 +16,25 @@ __all__ = ["synthesize_page", "synthesize_kie_page"] +class _Word(NamedTuple): + """A word to render: its text, where it starts, and how much room it was detected in.""" + + value: str + x: int # anchor: middle of the leading edge of the box, like the "lm" anchor of flat text + y: int + width: int # extent along the text direction + height: int # extent across the text direction + angle: float # degrees, counter-clockwise + + @lru_cache(maxsize=256) def _cached_font(font_family: str | None, font_size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: """Memoized font loader: avoids re-reading the font file for every word.""" - return get_font(font_family, font_size) + try: + return get_font(font_family, max(font_size, 1)) + except Exception: # noqa: BLE001 - a missing or broken font must not abort a whole page + logging.warning(f"Could not load font '{font_family}', falling back to the default font") + return get_font(None, max(font_size, 1)) @lru_cache(maxsize=1) @@ -28,67 +43,91 @@ def _warn_rotation_once() -> None: # pragma: no cover logging.warning("Polygons with larger rotations may lead to slightly inaccurate rendering") +def _points(geometry: Any) -> list[tuple[float, float]] | None: + try: + points = [(float(x), float(y)) for x, y in geometry] + except (TypeError, ValueError): + return None + if len(points) not in (2, 4) or not all(math.isfinite(v) for point in points for v in point): + return None + return points + + def _polygon_angle(polygon: list[tuple[float, float]], w: int, h: int) -> float: - """Estimate the rotation angle (degrees, counter-clockwise) from the top edge of a 4-point polygon.""" (x0, y0), (x1, y1) = polygon[0], polygon[1] return -math.degrees(math.atan2((y1 - y0) * h, (x1 - x0) * w)) def _text_width(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) -> int: bbox = font.getbbox(text) - return max(int(bbox[2]) - int(bbox[0]), 1) + return max(math.ceil(bbox[2]), 1) + + +def _text_height(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) -> int: + bbox = font.getbbox(text) + return max(int(bbox[3]) - int(bbox[1]), 1) + + +def _text_vspan(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) -> int: + """Ascender-to-descender span: the vertical extent the "lm" anchor is centered on.""" + try: + ascent, descent = font.getmetrics() + return max(ascent + descent, 1) + except AttributeError: # pragma: no cover - bitmap fonts expose no metrics + return _text_height(font, text) -def _fit_font( +def _fit_font_size( text: str, box_w: int, box_h: int, font_family: str | None, min_font_size: int, max_font_size: int, -) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: +) -> int: """Directly estimate the largest font size fitting the box (text width scales ~linearly with size).""" font_size = max(min(box_h, max_font_size), min_font_size) try: font = _cached_font(font_family, font_size) - x0, y0, x1, y1 = font.getbbox(text) - text_w, text_h = max(int(x1) - int(x0), 1), max(int(y1) - int(y0), 1) - if text_w > box_w or text_h > box_h: - scale = min(box_w / text_w, box_h / text_h) - font_size = max(min(int(font_size * scale), max_font_size), min_font_size) - font = _cached_font(font_family, font_size) - # The linear estimate can be off by a pixel or two: shrink until the text truly fits - while font_size > min_font_size and _text_width(font, text) > box_w: - font_size -= 1 - font = _cached_font(font_family, font_size) + # A font size is an em size while the box bounds the ink, so scale on the measured ink: + # without this the text is rendered noticeably smaller than the box it was detected in. + scale = min(box_w / _text_width(font, text), box_h / _text_height(font, text)) + font_size = max(min(int(font_size * scale), max_font_size), min_font_size) except ValueError: # pragma: no cover - font = _cached_font(font_family, min_font_size) - return font + font_size = min_font_size + return font_size -def _fit_line_font( - word_widths: list[tuple[str, int]], - line_height: int, +def _fit_line_font_size( + words: list[_Word], font_family: str | None, min_font_size: int, max_font_size: int, -) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: - """Find one font size for a whole line such that every word fits its own available width.""" - font_size = max(min(line_height, max_font_size), min_font_size) +) -> int: + """Find one font size for a whole line, driven by the median word box.""" + font_size = max(min(int(np.median([word.height for word in words])), max_font_size), min_font_size) try: font = _cached_font(font_family, font_size) - # Scale down so the most constrained word still fits its own box (linear estimate) - scale = min([avail_w / _text_width(font, value) for value, avail_w in word_widths] + [1.0]) - if scale < 1.0: - font_size = max(min(int(font_size * scale), max_font_size), min_font_size) - font = _cached_font(font_family, font_size) - # The linear estimate can be off by a pixel or two: shrink until every word truly fits - while font_size > min_font_size and any(_text_width(font, value) > avail_w for value, avail_w in word_widths): - font_size -= 1 - font = _cached_font(font_family, font_size) + scale = min( + float(np.median([word.height / _text_height(font, word.value) for word in words])), + float(np.median([word.width / _text_width(font, word.value) for word in words])), + ) + font_size = max(min(int(font_size * scale), max_font_size), min_font_size) except ValueError: # pragma: no cover - font = _cached_font(font_family, min_font_size) - return font + font_size = min_font_size + return font_size + + +def _harmonize(sizes: list[int], tolerance: float = 0.2) -> list[int]: + """Snap font sizes that are close to each other onto one common size. + + Boxes for one and the same typeface come out a pixel or two apart from line to line, and + rendering every line at its own size is what makes a synthesized page look ragged. Sizes + within `tolerance` of each other collapse onto their median, while genuinely different sizes + - a headline, a caption, small print - stay apart. + """ + values = np.array(sizes, dtype=float) + return [int(round(float(np.median(values[np.abs(values - size) <= tolerance * size])))) for size in sizes] def _draw_word( @@ -105,132 +144,311 @@ def _draw_word( except UnicodeEncodeError: d.text(xy, anyascii(text), font=font, fill=fill, anchor=anchor) except Exception: # pragma: no cover - logging.warning(f"Could not render word: {text}") + try: + # Anchors are rejected by bitmap fonts, which would otherwise leave the page blank + d.text(xy, anyascii(text), font=font, fill=fill) + except Exception: + logging.warning(f"Could not render word: {text}") -def _paste_rotated_word( +def _paste_word( response: Image.Image, text: str, font: ImageFont.FreeTypeFont | ImageFont.ImageFont, - center: tuple[int, int], - angle: float, + position: tuple[int, int], fill: tuple[int, int, int], + angle: float = 0.0, + squeeze: float = 1.0, ) -> None: - """Render a word on a transparent patch, rotate it, and paste it centered on the polygon centroid.""" - bbox = font.getbbox(text) - x0, y0, x1, y1 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) - patch = Image.new("RGBA", (max(x1 - x0, 1) + 4, max(y1 - y0, 1) + 4), (0, 0, 0, 0)) - _draw_word(ImageDraw.Draw(patch), (2 - x0, 2 - y0), text, font, fill, anchor="la") - patch = patch.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC) - response.paste(patch, (center[0] - patch.width // 2, center[1] - patch.height // 2), patch) + """Render a word on a transparent patch, condense and rotate it, then paste it. + + `position` is the middle of the left edge of the text, the same anchor point as a flat "lm" + draw. `squeeze` condenses the glyphs horizontally without touching their height, which is how + a line that is too long for its box stays inside it at the size the rest of the page uses. + """ + pad = 2 + patch = Image.new("RGBA", (_text_width(font, text) + 2 * pad, _text_vspan(font, text) + 2 * pad), (0, 0, 0, 0)) + _draw_word(ImageDraw.Draw(patch), (pad, pad), text, font, fill, anchor="la") + if squeeze < 1.0: + patch = patch.resize((max(round(patch.width * squeeze), 1), patch.height), Image.Resampling.BICUBIC) + pad = round(pad * squeeze) + + anchor_x, anchor_y = float(pad), patch.height / 2 + if abs(angle) > 3: + # rotate() turns the patch about its centre and expands the canvas, so follow the anchor + theta = math.radians(angle) + dx, dy = anchor_x - patch.width / 2, anchor_y - patch.height / 2 + patch = patch.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC) + anchor_x = patch.width / 2 + dx * math.cos(theta) + dy * math.sin(theta) + anchor_y = patch.height / 2 - dx * math.sin(theta) + dy * math.cos(theta) + response.paste(patch, (round(position[0] - anchor_x), round(position[1] - anchor_y)), patch) + + +def _entry_words(entry: dict[str, Any], w: int, h: int) -> list[_Word]: + """Collect the renderable words of a line entry, in pixel coordinates.""" + words = [] + for word in entry.get("words") or []: + if not isinstance(word, dict): + continue + geom = _points(word.get("geometry")) + if geom is None or not str(word.get("value", "")).strip(): + continue + value = str(word["value"]) + if len(geom) == 2: + (gx0, gy0), (gx1, gy1) = geom + x0, y0 = int(round(w * gx0)), int(round(h * gy0)) + x1, y1 = int(round(w * gx1)), int(round(h * gy1)) + words.append(_Word(value, min(x0, x1), (y0 + y1) // 2, max(abs(x1 - x0), 1), max(abs(y1 - y0), 1), 0.0)) + else: + # True text-direction extent: length of the top edge / left edge in pixels + width = math.hypot((geom[1][0] - geom[0][0]) * w, (geom[1][1] - geom[0][1]) * h) + height = math.hypot((geom[2][0] - geom[1][0]) * w, (geom[2][1] - geom[1][1]) * h) + words.append( + _Word( + value, + int(round(w * (geom[0][0] + geom[3][0]) / 2)), + int(round(h * (geom[0][1] + geom[3][1]) / 2)), + max(int(round(width)), 1), + max(int(round(height)), 1), + _polygon_angle(geom, w, h), + ) + ) + return words -def _synthesize( +def _line_axis(words: list[_Word]) -> tuple[float, float, float, float, float]: + """The baseline of a line: its origin, its unit direction and its angle. + + Words are placed along this one axis instead of at their own box centres, which is what keeps + a line straight; the perpendicular median absorbs boxes that sit a little high or low. + """ + angle = float(np.median([word.angle for word in words])) + theta = math.radians(angle) + ux, uy = math.cos(theta), -math.sin(theta) + x0, y0 = words[0].x, words[0].y + across = float(np.median([(word.x - x0) * -uy + (word.y - y0) * ux for word in words])) + return x0 + across * -uy, y0 + across * ux, ux, uy, angle + + +def _place_words(offsets: list[float], widths: list[float], space: float, budget: float) -> tuple[list[float], float]: + """Keep every word where it was detected, unless that would run it into its neighbour. + + A word only moves when the word before it needs the room, and if the line then overruns the + space it was detected in, the whole line is condensed by a single factor - so words never + overlap, and the ones that had to give way still share the size of the rest of the page. + """ + + def run(squeeze: float) -> list[float]: + cursor, placed = -math.inf, [] + for offset, width in zip(offsets, widths): + start = max(offset, cursor) + placed.append(start) + cursor = start + squeeze * (width + space) + return placed + + placed, squeeze = run(1.0), 1.0 + for _ in range(3): # pushes shrink as the line condenses, so a few passes converge + end = placed[-1] + squeeze * widths[-1] + if end <= budget or budget <= 0: + break + squeeze = max(squeeze * budget / end, 0.5) + placed = run(squeeze) + return placed, squeeze + + +def _synthesize_line( + response: Image.Image, + words: list[_Word], + font_size: int, + font_family: str | None, + min_font_size: int, + text_color: tuple[int, int, int], +) -> None: + """Draw the words of one line at one font size, spaced so that none can overlap the next.""" + ox, oy, ux, uy, angle = _line_axis(words) + along = sorted(((word.x - ox) * ux + (word.y - oy) * uy, word) for word in words) + offsets = [offset for offset, _ in along] + words = [word for _, word in along] + budget = max(offset + word.width for offset, word in along) + + for _ in range(2): + font = _cached_font(font_family, font_size) + squeezes = [max(min(1.0, word.width / _text_width(font, word.value)), 0.5) for word in words] + widths = [squeeze * _text_width(font, word.value) for squeeze, word in zip(squeezes, words)] + placed, line_squeeze = _place_words(offsets, widths, 1.0, budget) + # A line that would have to be condensed to less than half is not crowded but mis-sized: + # a smaller face keeps it readable instead of squashing the glyphs to nothing. + if line_squeeze > 0.5 or font_size <= min_font_size: + break + font_size = max(int(font_size * 2 * line_squeeze), min_font_size) + + d = ImageDraw.Draw(response) + for word, offset, squeeze in zip(words, placed, squeezes): + x, y = round(ox + offset * ux), round(oy + offset * uy) + squeeze *= line_squeeze + if abs(angle) > 3 or squeeze < 1.0: + _paste_word(response, word.value, font, (x, y), text_color, angle, squeeze) + else: + # "lm" anchor: vertically centered on the baseline, no ascender-offset drift + _draw_word(d, (x, y), word.value, font, text_color, anchor="lm") + + +def _synthesize_value( response: Image.Image, entry: dict[str, Any], + polygon: list[tuple[float, float]], + angle: float, w: int, h: int, - draw_proba: bool = False, - font_family: str | None = None, - min_font_size: int = 6, - max_font_size: int = 50, - text_color: tuple[int, int, int] = (0, 0, 0), -) -> Image.Image: - if len(entry["geometry"]) == 2: - (xmin_r, ymin_r), (xmax_r, ymax_r) = entry["geometry"] + font_size: int, + font_family: str | None, + text_color: tuple[int, int, int], +) -> None: + """Draw a single value (a word entry, or a KIE prediction) inside its own box.""" + text = str(entry["value"]) + # Measure along the text direction, not on the axis-aligned bounding box + box_w = max(int(round(math.hypot((polygon[1][0] - polygon[0][0]) * w, (polygon[1][1] - polygon[0][1]) * h))), 1) + font = _cached_font(font_family, font_size) + # Anchor on the middle of the leading edge, like the "lm" anchor of flat text + x = int(round(w * (polygon[0][0] + polygon[3][0]) / 2)) + y = int(round(h * (polygon[0][1] + polygon[3][1]) / 2)) + squeeze = min(1.0, box_w / _text_width(font, text)) # stay inside the box, do not run into the next field + if abs(angle) > 3 or squeeze < 1.0: + _paste_word(response, text, font, (x, y), text_color, angle, squeeze) + else: + _draw_word(ImageDraw.Draw(response), (x, y), text, font, text_color, anchor="lm") + + +def _draw_confidence( + response: Image.Image, + entry: dict[str, Any], + box: tuple[int, int, int, int], + font_family: str | None, +) -> None: + xmin, ymin, xmax, ymax = box + box_width, box_height = max(xmax - xmin, 1), max(ymax - ymin, 1) + confidences = [ + float(word["confidence"]) + for word in entry.get("words") or [] + if isinstance(word, dict) + and isinstance(word.get("confidence"), int | float) + and math.isfinite(float(word["confidence"])) + ] + confidence = entry.get("confidence") + if not (isinstance(confidence, int | float) and math.isfinite(float(confidence))): + confidence = float(np.mean(confidences)) if confidences else 1.0 + confidence = min(max(float(confidence), 0.0), 1.0) + p = int(255 * confidence) + color = (255 - p, 0, p) + + d = ImageDraw.Draw(response) + d.rectangle([(xmin, ymin), (xmax, ymax)], outline=color, width=2) + # Scale the confidence label with the box instead of a hardcoded size + prob_font = _cached_font(font_family, max(min(box_height // 2, 20), 10)) + prob_text = f"{confidence:.2f}" + prob_text_width, prob_text_height = prob_font.getbbox(prob_text)[2:4] + prob_x_offset = (box_width - prob_text_width) // 2 + prob_y_offset = max(0, ymin - prob_text_height - 2) + _draw_word(d, (int(xmin + prob_x_offset), int(prob_y_offset)), prob_text, prob_font, color, anchor="lt") + + +def _entry_geometry( + entry: dict[str, Any], w: int, h: int +) -> tuple[list[tuple[float, float]], float, tuple[int, int, int, int]] | None: + geometry = _points(entry.get("geometry")) + if geometry is None: + return None + if len(geometry) == 2: + (xmin_r, ymin_r), (xmax_r, ymax_r) = geometry polygon = [(xmin_r, ymin_r), (xmax_r, ymin_r), (xmax_r, ymax_r), (xmin_r, ymax_r)] angle = 0.0 else: - polygon = entry["geometry"] - angle = _polygon_angle(polygon, w, h) - - # Calculate the bounding box of the entry + polygon, angle = geometry, _polygon_angle(geometry, w, h) + _warn_rotation_once() # pragma: no cover x_coords, y_coords = zip(*polygon) - xmin, ymin, xmax, ymax = ( + box = ( int(round(w * min(x_coords))), int(round(h * min(y_coords))), int(round(w * max(x_coords))), int(round(h * max(y_coords))), ) - box_width, box_height = max(xmax - xmin, 1), max(ymax - ymin, 1) + return polygon, angle, box - d = ImageDraw.Draw(response) - if "words" in entry: - # Line entry: one consistent font size for the whole line, drawn word by word. - word_render: list[tuple[str, int, int, int, int, float]] = [] - for word in entry["words"]: - geom = word["geometry"] - if len(geom) == 2: - (gx0, gy0), (gx1, gy1) = geom - wxmin, wymin = int(round(w * gx0)), int(round(h * gy0)) - wxmax, wymax = int(round(w * gx1)), int(round(h * gy1)) - word_render.append(( - word["value"], - wxmin, - (wymin + wymax) // 2, - max(wxmax - wxmin, 1), - max(wymax - wymin, 1), - 0.0, - )) - else: - xs, ys = zip(*geom) - cx = int(round(w * sum(xs) / len(xs))) - cy = int(round(h * sum(ys) / len(ys))) - # True text-direction extent: length of the top edge / left edge in pixels - avail_w = int(round(math.hypot((geom[1][0] - geom[0][0]) * w, (geom[1][1] - geom[0][1]) * h))) - avail_h = int(round(math.hypot((geom[2][0] - geom[1][0]) * w, (geom[2][1] - geom[1][1]) * h))) - word_render.append(( - word["value"], - cx, - cy, - max(avail_w, 1), - max(avail_h, 1), - _polygon_angle(geom, w, h), - )) - line_height = min(avail_h for *_, avail_h, _angle in word_render) - font = _fit_line_font( - [(value, avail_w) for value, _, _, avail_w, _, _ in word_render], - line_height, - font_family, - min_font_size, - max_font_size, - ) - for value, ax, ay, _, _, word_angle in word_render: - if abs(word_angle) > 3: - _paste_rotated_word(response, value, font, (ax, ay), word_angle, text_color) - else: - _draw_word(d, (ax, ay), value, font, text_color, anchor="lm") - else: - word_text = entry["value"] - if abs(angle) > 3: # Rotated word: render on a patch and paste it rotated - font = _fit_font(word_text, box_width, box_height, font_family, min_font_size, max_font_size) - cx, cy = int(round(w * sum(x_coords) / len(x_coords))), int(round(h * sum(y_coords) / len(y_coords))) - _paste_rotated_word(response, word_text, font, (cx, cy), angle, text_color) - else: - font = _fit_font(word_text, box_width, box_height, font_family, min_font_size, max_font_size) - # "lm" anchor: vertically centered in the box, no ascender-offset drift - _draw_word(d, (xmin, (ymin + ymax) // 2), word_text, font, text_color, anchor="lm") - - if draw_proba: - confidence = ( - entry["confidence"] - if "confidence" in entry - else sum(word["confidence"] for word in entry["words"]) / len(entry["words"]) - ) - p = int(255 * confidence) - color = (255 - p, 0, p) # Red to blue gradient based on probability - d.rectangle([(xmin, ymin), (xmax, ymax)], outline=color, width=2) +def _entry_font_size( + entry: dict[str, Any], + words: list[_Word], + polygon: list[tuple[float, float]], + w: int, + h: int, + font_family: str | None, + min_font_size: int, + max_font_size: int, +) -> int: + if words: + return _fit_line_font_size(words, font_family, min_font_size, max_font_size) + box_w = max(int(round(math.hypot((polygon[1][0] - polygon[0][0]) * w, (polygon[1][1] - polygon[0][1]) * h))), 1) + box_h = max(int(round(math.hypot((polygon[2][0] - polygon[1][0]) * w, (polygon[2][1] - polygon[1][1]) * h))), 1) + return _fit_font_size(str(entry["value"]), box_w, box_h, font_family, min_font_size, max_font_size) - # Scale the confidence label with the box instead of a hardcoded size - prob_font = _cached_font(font_family, max(min(box_height // 2, 20), 10)) - prob_text = f"{confidence:.2f}" - prob_text_width, prob_text_height = prob_font.getbbox(prob_text)[2:4] - prob_x_offset = (box_width - prob_text_width) // 2 - prob_y_offset = max(0, ymin - prob_text_height - 2) - d.text((xmin + prob_x_offset, prob_y_offset), prob_text, font=prob_font, fill=color, anchor="lt") - return response +def _page_size(page: dict[str, Any]) -> tuple[int, int]: + try: + h, w = (int(round(float(value))) for value in page["dimensions"]) + except Exception as exc: + raise ValueError(f"invalid page 'dimensions': {page.get('dimensions')!r}") from exc + if h <= 0 or w <= 0: + raise ValueError(f"page dimensions must be positive, got {(h, w)}") + if Image.MAX_IMAGE_PIXELS is not None and h * w > Image.MAX_IMAGE_PIXELS: + raise ValueError(f"page dimensions {(h, w)} exceed Pillow's MAX_IMAGE_PIXELS safety limit") + return h, w + + +def _render( + page: dict[str, Any], + entries: list[dict[str, Any]], + draw_proba: bool, + font_family: str | None, + min_font_size: int, + max_font_size: int, + background_color: tuple[int, int, int], + text_color: tuple[int, int, int], +) -> np.ndarray: + """Render entries onto a blank page at one harmonized set of font sizes.""" + h, w = _page_size(page) + response = Image.new("RGB", (w, h), color=background_color) + + prepared = [] + for entry in entries: + geometry = _entry_geometry(entry, w, h) + if geometry is None: + continue + polygon, angle, box = geometry + words = _entry_words(entry, w, h) + if not words and not str(entry.get("value", "")).strip(): + if draw_proba: + _draw_confidence(response, entry, box, font_family) + continue + try: + size = _entry_font_size(entry, words, polygon, w, h, font_family, min_font_size, max_font_size) + except Exception as exc: # noqa: BLE001 + logging.warning(f"Could not size entry: {exc}") + continue + prepared.append((entry, polygon, angle, box, words, size)) + + sizes = _harmonize([size for *_, size in prepared]) if prepared else [] + for (entry, polygon, angle, box, words, _), size in zip(prepared, sizes): + try: + if words: + _synthesize_line(response, words, size, font_family, min_font_size, text_color) + else: + _synthesize_value(response, entry, polygon, angle, w, h, size, font_family, text_color) + except Exception as exc: # noqa: BLE001 + logging.warning(f"Could not render entry: {exc}") + if draw_proba: + _draw_confidence(response, entry, box, font_family) + + return np.array(response, dtype=np.uint8) def synthesize_page( @@ -256,27 +474,14 @@ def synthesize_page( Returns: the synthesized page """ - h, w = page["dimensions"] - response = Image.new("RGB", (w, h), color=background_color) - - for block in page["blocks"]: - for line in block["lines"]: - if len(line["geometry"]) == 4: - _warn_rotation_once() # pragma: no cover - # Line-level entry keeps a consistent font per line while preserving word positions - response = _synthesize( - response=response, - entry=line, - w=w, - h=h, - draw_proba=draw_proba, - font_family=font_family, - min_font_size=min_font_size, - max_font_size=max_font_size, - text_color=text_color, - ) - - return np.array(response, dtype=np.uint8) + lines = [ + line + for block in page.get("blocks") or [] + if isinstance(block, dict) + for line in block.get("lines") or [] + if isinstance(line, dict) + ] + return _render(page, lines, draw_proba, font_family, min_font_size, max_font_size, background_color, text_color) def synthesize_kie_page( @@ -302,22 +507,12 @@ def synthesize_kie_page( Returns: the synthesized page """ - h, w = page["dimensions"] - response = Image.new("RGB", (w, h), color=background_color) - - for predictions in page["predictions"].values(): - for prediction in predictions: - if len(prediction["geometry"]) == 4: - _warn_rotation_once() # pragma: no cover - response = _synthesize( - response=response, - entry=prediction, - w=w, - h=h, - draw_proba=draw_proba, - font_family=font_family, - min_font_size=min_font_size, - max_font_size=max_font_size, - text_color=text_color, - ) - return np.array(response, dtype=np.uint8) + predictions = [ + prediction + for predictions in (page.get("predictions") or {}).values() + for prediction in predictions or [] + if isinstance(prediction, dict) + ] + return _render( + page, predictions, draw_proba, font_family, min_font_size, max_font_size, background_color, text_color + ) From 39deebe5441c37ed49c724b6a61bceaba0babce1 Mon Sep 17 00:00:00 2001 From: felix Date: Fri, 31 Jul 2026 09:46:55 +0200 Subject: [PATCH 2/3] improve --- doctr/utils/reconstitution.py | 126 ++++++++++++++++++--- tests/common/test_utils_reconstitution.py | 130 ++++++++++++++++++++-- 2 files changed, 235 insertions(+), 21 deletions(-) diff --git a/doctr/utils/reconstitution.py b/doctr/utils/reconstitution.py index acc734e976..28f486914e 100644 --- a/doctr/utils/reconstitution.py +++ b/doctr/utils/reconstitution.py @@ -32,7 +32,7 @@ def _cached_font(font_family: str | None, font_size: int) -> ImageFont.FreeTypeF """Memoized font loader: avoids re-reading the font file for every word.""" try: return get_font(font_family, max(font_size, 1)) - except Exception: # noqa: BLE001 - a missing or broken font must not abort a whole page + except Exception: # pragma: no cover logging.warning(f"Could not load font '{font_family}', falling back to the default font") return get_font(None, max(font_size, 1)) @@ -44,6 +44,7 @@ def _warn_rotation_once() -> None: # pragma: no cover def _points(geometry: Any) -> list[tuple[float, float]] | None: + """Validate a geometry and return its points, or None if it is unusable.""" try: points = [(float(x), float(y)) for x, y in geometry] except (TypeError, ValueError): @@ -54,16 +55,19 @@ def _points(geometry: Any) -> list[tuple[float, float]] | None: def _polygon_angle(polygon: list[tuple[float, float]], w: int, h: int) -> float: + """Estimate the rotation angle (degrees, counter-clockwise) from the top edge of a 4-point polygon.""" (x0, y0), (x1, y1) = polygon[0], polygon[1] return -math.degrees(math.atan2((y1 - y0) * h, (x1 - x0) * w)) def _text_width(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) -> int: + """Width from the drawing origin to the right edge of the ink, so the left bearing counts.""" bbox = font.getbbox(text) return max(math.ceil(bbox[2]), 1) def _text_height(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) -> int: + """Height from the top of the ink to the bottom, so the ascender counts.""" bbox = font.getbbox(text) return max(int(bbox[3]) - int(bbox[1]), 1) @@ -71,7 +75,7 @@ def _text_height(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) def _text_vspan(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) -> int: """Ascender-to-descender span: the vertical extent the "lm" anchor is centered on.""" try: - ascent, descent = font.getmetrics() + ascent, descent = font.getmetrics() # type: ignore[union-attr] return max(ascent + descent, 1) except AttributeError: # pragma: no cover - bitmap fonts expose no metrics return _text_height(font, text) @@ -104,7 +108,13 @@ def _fit_line_font_size( min_font_size: int, max_font_size: int, ) -> int: - """Find one font size for a whole line, driven by the median word box.""" + """Find one font size for a whole line, driven by the median word box. + + Using the median rather than the smallest box is what keeps a single tiny box - a full stop, + a subscript, a mis-detected accent - from collapsing the entire line. The width is measured + the same way: a summed width would let the slack of a generous box pay for the overflow of a + tight one, which sizes the line too large for half of its words. + """ font_size = max(min(int(np.median([word.height for word in words])), max_font_size), min_font_size) try: font = _cached_font(font_family, font_size) @@ -138,6 +148,7 @@ def _draw_word( fill: tuple[int, int, int], anchor: str = "lm", ) -> None: + """Draw a word, falling back to ASCII if the font cannot render it.""" try: try: d.text(xy, text, font=font, fill=fill, anchor=anchor) @@ -184,6 +195,48 @@ def _paste_word( response.paste(patch, (round(position[0] - anchor_x), round(position[1] - anchor_y)), patch) +def _tilt(words: list[_Word]) -> float | None: + """Direction of a line in degrees, read from where its words sit, or None if they disagree.""" + if len(words) < 4: + return None + xs = np.array([word.x for word in words], dtype=float) + ys = np.array([word.y for word in words], dtype=float) + height = float(np.median([word.height for word in words])) + dx, dy = xs[:, None] - xs[None, :], ys[:, None] - ys[None, :] + spread = np.abs(dx) > height # pairs too close together only measure box noise + if not spread.any(): + return None + slopes = dy[spread] / dx[spread] + slope = float(np.median(slopes)) + if float(np.median(np.abs(slopes - slope))) > max(0.5 * abs(slope), 0.02): + return None + return -math.degrees(math.atan(slope)) + + +def _unrotate_box(word: _Word, angle: float) -> _Word: + """Recover the text extent and leading-edge anchor of a word from its axis-aligned box.""" + s, c = abs(math.sin(math.radians(angle))), abs(math.cos(math.radians(angle))) + det = c * c - s * s + if det < 0.2: # past ~39 degrees the inversion stops being trustworthy + return word + width = (word.width * c - word.height * s) / det + height = (word.height * c - word.width * s) / det + if width < 1 or height < 1: + return word + # The anchor moves from the middle of the box edge to the middle of the leading edge of the text + x = word.x + height * s / 2 + y = word.y - word.height / 2 + height * c / 2 if angle < 0 else word.y + word.height / 2 - height * c / 2 + return word._replace(x=round(x), y=round(y), width=round(width), height=round(height), angle=angle) + + +def _space_width(font: ImageFont.FreeTypeFont | ImageFont.ImageFont) -> float: + """Half of the font's own space advance: the least gap that still reads as a word break.""" + try: + return max(float(font.getlength(" ")) / 2, 1.0) + except Exception: # noqa: BLE001 - pragma: no cover - a bitmap font may not measure a space + return max(_text_width(font, "n") / 2, 1.0) + + def _entry_words(entry: dict[str, Any], w: int, h: int) -> list[_Word]: """Collect the renderable words of a line entry, in pixel coordinates.""" words = [] @@ -213,16 +266,19 @@ def _entry_words(entry: dict[str, Any], w: int, h: int) -> list[_Word]: _polygon_angle(geom, w, h), ) ) - return words + # An upright box carries no angle of its own, so a tilted line has to be recognised from where + # its words sit before anything can be sized off boxes that the tilt has inflated + tilt = _tilt(words) if words and not any(word.angle for word in words) else None + return [_unrotate_box(word, tilt) for word in words] if tilt and abs(tilt) >= 3 else words def _line_axis(words: list[_Word]) -> tuple[float, float, float, float, float]: - """The baseline of a line: its origin, its unit direction and its angle. - - Words are placed along this one axis instead of at their own box centres, which is what keeps - a line straight; the perpendicular median absorbs boxes that sit a little high or low. - """ - angle = float(np.median([word.angle for word in words])) + """The baseline of a line: its origin, its unit direction and its angle.""" + # Where the words sit pins the direction down far more tightly than their own corners do, + # so the box angles are only a fallback for lines too short to read a trend from + angle = _tilt(words) + if angle is None: + angle = _median_angle(words) theta = math.radians(angle) ux, uy = math.cos(theta), -math.sin(theta) x0, y0 = words[0].x, words[0].y @@ -256,6 +312,29 @@ def run(squeeze: float) -> list[float]: return placed, squeeze +def _median_angle(words: list[_Word]) -> float: + """Median of the word angles, folded so they all lie within a quarter turn of zero.""" + return float(np.median([(word.angle + 90) % 180 - 90 for word in words])) + + +def _split_rows(words: list[_Word]) -> list[list[_Word]]: + """Split a group of words into the separate baselines they actually sit on.""" + if len(words) < 2: + return [words] + theta = math.radians(_median_angle(words)) + vx, vy = math.sin(theta), math.cos(theta) # across the text direction + tolerance = 0.8 * float(np.median([word.height for word in words])) + ordered = sorted(words, key=lambda word: word.x * vx + word.y * vy) + rows, current = [], [ordered[0]] + for previous, word in zip(ordered, ordered[1:]): + if (word.x - previous.x) * vx + (word.y - previous.y) * vy > tolerance: + rows.append(current) + current = [] + current.append(word) + rows.append(current) + return rows + + def _synthesize_line( response: Image.Image, words: list[_Word], @@ -264,7 +343,20 @@ def _synthesize_line( min_font_size: int, text_color: tuple[int, int, int], ) -> None: - """Draw the words of one line at one font size, spaced so that none can overlap the next.""" + """Draw a line, as one row per baseline its words turn out to share.""" + for row in _split_rows(words): + _synthesize_row(response, row, font_size, font_family, min_font_size, text_color) + + +def _synthesize_row( + response: Image.Image, + words: list[_Word], + font_size: int, + font_family: str | None, + min_font_size: int, + text_color: tuple[int, int, int], +) -> None: + """Draw the words of one row at one font size, spaced so that none can overlap the next.""" ox, oy, ux, uy, angle = _line_axis(words) along = sorted(((word.x - ox) * ux + (word.y - oy) * uy, word) for word in words) offsets = [offset for offset, _ in along] @@ -275,7 +367,9 @@ def _synthesize_line( font = _cached_font(font_family, font_size) squeezes = [max(min(1.0, word.width / _text_width(font, word.value)), 0.5) for word in words] widths = [squeeze * _text_width(font, word.value) for squeeze, word in zip(squeezes, words)] - placed, line_squeeze = _place_words(offsets, widths, 1.0, budget) + # The boxes cannot supply the spacing (a detector dilates them), so the layout keeps at + # least a readable gap of its own between one word and the next + placed, line_squeeze = _place_words(offsets, widths, _space_width(font), budget) # A line that would have to be condensed to less than half is not crowded but mis-sized: # a smaller face keeps it readable instead of squashing the glyphs to nothing. if line_squeeze > 0.5 or font_size <= min_font_size: @@ -325,6 +419,7 @@ def _draw_confidence( box: tuple[int, int, int, int], font_family: str | None, ) -> None: + """Outline the entry and label it with its confidence. Blue: p=1, red: p=0.""" xmin, ymin, xmax, ymax = box box_width, box_height = max(xmax - xmin, 1), max(ymax - ymin, 1) confidences = [ @@ -355,6 +450,7 @@ def _draw_confidence( def _entry_geometry( entry: dict[str, Any], w: int, h: int ) -> tuple[list[tuple[float, float]], float, tuple[int, int, int, int]] | None: + """Normalize an entry geometry to a 4-point polygon, its angle and its pixel bounding box.""" geometry = _points(entry.get("geometry")) if geometry is None: return None @@ -385,6 +481,7 @@ def _entry_font_size( min_font_size: int, max_font_size: int, ) -> int: + """The size this entry would like, before it is harmonized with the rest of the page.""" if words: return _fit_line_font_size(words, font_family, min_font_size, max_font_size) box_w = max(int(round(math.hypot((polygon[1][0] - polygon[0][0]) * w, (polygon[1][1] - polygon[0][1]) * h))), 1) @@ -393,6 +490,7 @@ def _entry_font_size( def _page_size(page: dict[str, Any]) -> tuple[int, int]: + """Validate the page dimensions before allocating the canvas.""" try: h, w = (int(round(float(value))) for value in page["dimensions"]) except Exception as exc: @@ -431,7 +529,7 @@ def _render( continue try: size = _entry_font_size(entry, words, polygon, w, h, font_family, min_font_size, max_font_size) - except Exception as exc: # noqa: BLE001 + except Exception as exc: logging.warning(f"Could not size entry: {exc}") continue prepared.append((entry, polygon, angle, box, words, size)) @@ -443,7 +541,7 @@ def _render( _synthesize_line(response, words, size, font_family, min_font_size, text_color) else: _synthesize_value(response, entry, polygon, angle, w, h, size, font_family, text_color) - except Exception as exc: # noqa: BLE001 + except Exception as exc: logging.warning(f"Could not render entry: {exc}") if draw_proba: _draw_confidence(response, entry, box, font_family) diff --git a/tests/common/test_utils_reconstitution.py b/tests/common/test_utils_reconstitution.py index 3632870998..2afb52dca5 100644 --- a/tests/common/test_utils_reconstitution.py +++ b/tests/common/test_utils_reconstitution.py @@ -1,3 +1,5 @@ +import math + import numpy as np from test_io_elements import _mock_kie_pages, _mock_pages @@ -104,8 +106,6 @@ def test_synthesize_kie_page_rotated_prediction(caplog): def test_synthesize_page_words_do_not_overlap(): - # Two adjacent words whose boxes are much narrower than the naive line-level font - # would require: the render must keep the gap between the boxes blank (regression test) page = { "dimensions": (300, 400), "blocks": [ @@ -132,11 +132,6 @@ def test_synthesize_page_words_do_not_overlap(): def test_synthesize_page_rotated_line(): - # A line whose words carry rotated 4-point polygons must render without error, - # follow the rotation (ink appears along the tilted baseline, not just the top band), - # and keep adjacent words from overlapping - import math - h_px, w_px = 400, 600 angle = math.radians(-18) dx, dy = math.cos(angle), math.sin(angle) @@ -174,3 +169,124 @@ def rot_word(value, start_x, start_y, width): # a horizontal per-bbox render would not place ink that high at those x-positions right_half = render[: 220 - 2 * height, w_px // 2 :] assert (right_half < 128).any() + + +def test_synthesize_page_tilted_straight_boxes(): + h_px, w_px = 700, 1400 + tilt, height = math.radians(20), 40 + words, x = [], 100.0 + for token in ("Photographed", "a", "pages", "of", "tilted", "i", "documents", "to", "read"): + text_w = 26 * len(token) + 10 + y = 150.0 + (x - 100) * math.tan(tilt) + box_w = text_w * math.cos(tilt) + box_h = text_w * math.sin(tilt) + height * math.cos(tilt) + words.append({ + "value": token, + "confidence": 0.9, + "geometry": ((x / w_px, y / h_px), ((x + box_w) / w_px, (y + box_h) / h_px)), + }) + x += box_w + 18 + xs = [c for word in words for c in (word["geometry"][0][0], word["geometry"][1][0])] + ys = [c for word in words for c in (word["geometry"][0][1], word["geometry"][1][1])] + geometry = ((min(xs), min(ys)), (max(xs), max(ys))) + page = { + "dimensions": (h_px, w_px), + "blocks": [{"geometry": geometry, "lines": [{"geometry": geometry, "words": words}]}], + } + render = reconstitution.synthesize_page(page) + _assert_valid_render(render, (h_px, w_px)) + + # The ink must sit on one straight line: measured across its own direction it can only spread + # as far as the text is tall, not as far as the boxes were inflated + ink = np.argwhere(render.min(axis=2) < 128).astype(float) + centred = ink - ink.mean(axis=0) + _, _, axes = np.linalg.svd(centred, full_matrices=False) + assert np.abs(centred @ axes[1]).max() < 40 + + +def test_synthesize_page_words_stay_apart_when_boxes_touch(): + page = { + "dimensions": (300, 900), + "blocks": [ + { + "geometry": ((0.05, 0.35), (0.80, 0.62)), + "lines": [ + { + "geometry": ((0.05, 0.35), (0.80, 0.62)), + "words": [ + {"value": "dominant", "confidence": 0.9, "geometry": ((0.05, 0.35), (0.32, 0.62))}, + {"value": "language", "confidence": 0.9, "geometry": ((0.30, 0.35), (0.57, 0.62))}, + {"value": "used", "confidence": 0.9, "geometry": ((0.55, 0.35), (0.80, 0.62))}, + ], + } + ], + } + ], + } + render = reconstitution.synthesize_page(page) + _assert_valid_render(render, (300, 900)) + + # Blank columns split the line into the gaps between letters and the two wider ones between + # the three words - if the words run together the widest gaps are just letter gaps + columns = np.nonzero((render.min(axis=2) < 128).any(axis=0))[0] + blanks = sorted((b - a - 1 for a, b in zip(columns, columns[1:]) if b - a > 1), reverse=True) + assert blanks[1] > 1.5 * float(np.median(blanks)) + + +def test_synthesize_page_line_angle_survives_noisy_boxes(): + h_px, w_px = 600, 1600 + words, x = [], 80.0 + for idx, token in enumerate(("required", "a", "large", "cluster", "of", "boxes", "or", "access", "to", "read")): + box_w = 22 * len(token) + 12 + lean = 5 if idx % 4 else -5 # most boxes lean one way, some the other + corners = ((x, 300 + lean), (x + box_w, 300 - lean), (x + box_w, 334 - lean), (x, 334 + lean)) + words.append({ + "value": token, + "confidence": 0.9, + "geometry": [(px / w_px, py / h_px) for px, py in corners], + }) + x += box_w + 16 + xs = [p[0] for word in words for p in word["geometry"]] + ys = [p[1] for word in words for p in word["geometry"]] + geometry = ((min(xs), min(ys)), (max(xs), max(ys))) + page = { + "dimensions": (h_px, w_px), + "blocks": [{"geometry": geometry, "lines": [{"geometry": geometry, "words": words}]}], + } + render = reconstitution.synthesize_page(page) + _assert_valid_render(render, (h_px, w_px)) + + # The words all sit on a flat baseline, so the render must be flat too: a line drawn at the + # median of the leaning box angles climbs or falls by tens of pixels across the page + ink = np.argwhere(render.min(axis=2) < 128) + left, right = ink[ink[:, 1] < w_px // 3], ink[ink[:, 1] > 2 * w_px // 3] + assert abs(right[:, 0].mean() - left[:, 0].mean()) < 12 + + +def test_synthesize_page_line_holding_two_rows(): + h_px, w_px = 400, 800 + words = [] + for tokens, top in ((("first", "row", "of", "words"), 120), (("second", "row", "here"), 180)): + x = 60.0 + for token in tokens: + box_w = 16 * len(token) + words.append({ + "value": token, + "confidence": 0.9, + "geometry": ((x / w_px, top / h_px), ((x + box_w) / w_px, (top + 30) / h_px)), + }) + x += box_w + 14 + xs = [c for word in words for c in (word["geometry"][0][0], word["geometry"][1][0])] + ys = [c for word in words for c in (word["geometry"][0][1], word["geometry"][1][1])] + geometry = ((min(xs), min(ys)), (max(xs), max(ys))) + page = { + "dimensions": (h_px, w_px), + "blocks": [{"geometry": geometry, "lines": [{"geometry": geometry, "words": words}]}], + } + render = reconstitution.synthesize_page(page) + _assert_valid_render(render, (h_px, w_px)) + + # Both rows must carry ink and the band between them stay clear; flattening the line onto one + # baseline empties one row and fills the gap + assert all((render[top : top + 30].min(axis=2) < 128).sum() > 50 for top in (120, 180)) + assert (render[154:176] == 255).all() From 186d29b73ce830c562f447f6c430863c470c8b97 Mon Sep 17 00:00:00 2001 From: felix Date: Fri, 31 Jul 2026 15:02:41 +0200 Subject: [PATCH 3/3] Improve layout aware --- doctr/utils/reconstitution.py | 195 ++++++++++++++++++---- tests/common/test_utils_reconstitution.py | 171 ++++++++++++++++++- 2 files changed, 337 insertions(+), 29 deletions(-) diff --git a/doctr/utils/reconstitution.py b/doctr/utils/reconstitution.py index 28f486914e..70e4a9da9a 100644 --- a/doctr/utils/reconstitution.py +++ b/doctr/utils/reconstitution.py @@ -60,10 +60,10 @@ def _polygon_angle(polygon: list[tuple[float, float]], w: int, h: int) -> float: return -math.degrees(math.atan2((y1 - y0) * h, (x1 - x0) * w)) -def _text_width(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) -> int: +def _text_width(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str, stroke: int = 0) -> int: """Width from the drawing origin to the right edge of the ink, so the left bearing counts.""" bbox = font.getbbox(text) - return max(math.ceil(bbox[2]), 1) + return max(math.ceil(bbox[2]) + 2 * stroke, 1) def _text_height(font: ImageFont.FreeTypeFont | ImageFont.ImageFont, text: str) -> int: @@ -147,13 +147,18 @@ def _draw_word( font: ImageFont.FreeTypeFont | ImageFont.ImageFont, fill: tuple[int, int, int], anchor: str = "lm", + stroke: int = 0, ) -> None: - """Draw a word, falling back to ASCII if the font cannot render it.""" + """Draw a word, falling back to ASCII if the font cannot render it. + + `stroke` outlines the glyphs in their own colour, which is how a heading is set in bold without + a second font file: doctr resolves one family, and it is not always shipped with a bold face. + """ try: try: - d.text(xy, text, font=font, fill=fill, anchor=anchor) + d.text(xy, text, font=font, fill=fill, anchor=anchor, stroke_width=stroke, stroke_fill=fill) except UnicodeEncodeError: - d.text(xy, anyascii(text), font=font, fill=fill, anchor=anchor) + d.text(xy, anyascii(text), font=font, fill=fill, anchor=anchor, stroke_width=stroke, stroke_fill=fill) except Exception: # pragma: no cover try: # Anchors are rejected by bitmap fonts, which would otherwise leave the page blank @@ -170,6 +175,7 @@ def _paste_word( fill: tuple[int, int, int], angle: float = 0.0, squeeze: float = 1.0, + stroke: int = 0, ) -> None: """Render a word on a transparent patch, condense and rotate it, then paste it. @@ -177,9 +183,13 @@ def _paste_word( draw. `squeeze` condenses the glyphs horizontally without touching their height, which is how a line that is too long for its box stays inside it at the size the rest of the page uses. """ - pad = 2 - patch = Image.new("RGBA", (_text_width(font, text) + 2 * pad, _text_vspan(font, text) + 2 * pad), (0, 0, 0, 0)) - _draw_word(ImageDraw.Draw(patch), (pad, pad), text, font, fill, anchor="la") + pad = 2 + stroke + patch = Image.new( + "RGBA", + (_text_width(font, text, stroke) + 2 * pad, _text_vspan(font, text) + 2 * pad + 2 * stroke), + (0, 0, 0, 0), + ) + _draw_word(ImageDraw.Draw(patch), (pad, pad + stroke), text, font, fill, anchor="la", stroke=stroke) if squeeze < 1.0: patch = patch.resize((max(round(patch.width * squeeze), 1), patch.height), Image.Resampling.BICUBIC) pad = round(pad * squeeze) @@ -335,6 +345,25 @@ def _split_rows(words: list[_Word]) -> list[list[_Word]]: return rows +def _place_word( + response: Image.Image, + d: ImageDraw.ImageDraw, + text: str, + font: ImageFont.FreeTypeFont | ImageFont.ImageFont, + position: tuple[int, int], + fill: tuple[int, int, int], + angle: float = 0.0, + squeeze: float = 1.0, + stroke: int = 0, +) -> None: + """Put one word on the page: straight onto it when it can be, on a patch when it cannot.""" + if abs(angle) > 3 or squeeze < 1.0: + _paste_word(response, text, font, position, fill, angle, squeeze, stroke) + else: + # "lm" anchor: vertically centered on the baseline, no ascender-offset drift + _draw_word(d, position, text, font, fill, anchor="lm", stroke=stroke) + + def _synthesize_line( response: Image.Image, words: list[_Word], @@ -342,10 +371,11 @@ def _synthesize_line( font_family: str | None, min_font_size: int, text_color: tuple[int, int, int], + bold: bool = False, ) -> None: """Draw a line, as one row per baseline its words turn out to share.""" for row in _split_rows(words): - _synthesize_row(response, row, font_size, font_family, min_font_size, text_color) + _synthesize_row(response, row, font_size, font_family, min_font_size, text_color, bold) def _synthesize_row( @@ -355,6 +385,7 @@ def _synthesize_row( font_family: str | None, min_font_size: int, text_color: tuple[int, int, int], + bold: bool = False, ) -> None: """Draw the words of one row at one font size, spaced so that none can overlap the next.""" ox, oy, ux, uy, angle = _line_axis(words) @@ -365,8 +396,11 @@ def _synthesize_row( for _ in range(2): font = _cached_font(font_family, font_size) - squeezes = [max(min(1.0, word.width / _text_width(font, word.value)), 0.5) for word in words] - widths = [squeeze * _text_width(font, word.value) for squeeze, word in zip(squeezes, words)] + # A bold face is set by outlining the glyphs, which widens them: the layout has to allow + # for that or the extra weight would spill over the following word + stroke = max(round(font_size / 32), 1) if bold else 0 + squeezes = [max(min(1.0, word.width / _text_width(font, word.value, stroke)), 0.5) for word in words] + widths = [squeeze * _text_width(font, word.value, stroke) for squeeze, word in zip(squeezes, words)] # The boxes cannot supply the spacing (a detector dilates them), so the layout keeps at # least a readable gap of its own between one word and the next placed, line_squeeze = _place_words(offsets, widths, _space_width(font), budget) @@ -379,12 +413,7 @@ def _synthesize_row( d = ImageDraw.Draw(response) for word, offset, squeeze in zip(words, placed, squeezes): x, y = round(ox + offset * ux), round(oy + offset * uy) - squeeze *= line_squeeze - if abs(angle) > 3 or squeeze < 1.0: - _paste_word(response, word.value, font, (x, y), text_color, angle, squeeze) - else: - # "lm" anchor: vertically centered on the baseline, no ascender-offset drift - _draw_word(d, (x, y), word.value, font, text_color, anchor="lm") + _place_word(response, d, word.value, font, (x, y), text_color, angle, squeeze * line_squeeze, stroke) def _synthesize_value( @@ -397,6 +426,7 @@ def _synthesize_value( font_size: int, font_family: str | None, text_color: tuple[int, int, int], + bold: bool = False, ) -> None: """Draw a single value (a word entry, or a KIE prediction) inside its own box.""" text = str(entry["value"]) @@ -406,11 +436,10 @@ def _synthesize_value( # Anchor on the middle of the leading edge, like the "lm" anchor of flat text x = int(round(w * (polygon[0][0] + polygon[3][0]) / 2)) y = int(round(h * (polygon[0][1] + polygon[3][1]) / 2)) - squeeze = min(1.0, box_w / _text_width(font, text)) # stay inside the box, do not run into the next field - if abs(angle) > 3 or squeeze < 1.0: - _paste_word(response, text, font, (x, y), text_color, angle, squeeze) - else: - _draw_word(ImageDraw.Draw(response), (x, y), text, font, text_color, anchor="lm") + stroke = max(round(font_size / 64), 1) if bold else 0 + # stay inside the box, do not run into the next field + squeeze = min(1.0, box_w / _text_width(font, text, stroke)) + _place_word(response, ImageDraw.Draw(response), text, font, (x, y), text_color, angle, squeeze, stroke) def _draw_confidence( @@ -447,6 +476,37 @@ def _draw_confidence( _draw_word(d, (int(xmin + prob_x_offset), int(prob_y_offset)), prob_text, prob_font, color, anchor="lt") +def _region_kind(kind: str) -> str: + """Normalize a layout class name so 'Section-header', 'section header' and 'SECTION_HEADER' match.""" + return "".join(character for character in kind.lower() if character.isalnum()) + + +def _region_type(regions: list[tuple[tuple[int, int, int, int], str]], point: tuple[int, int]) -> str: + """Type of the smallest layout region holding a point, or "" if it falls outside them all.""" + holding = [(box, kind) for box, kind in regions if box[0] <= point[0] <= box[2] and box[1] <= point[1] <= box[3]] + if not holding: + return "" + box, kind = min(holding, key=lambda item: (item[0][2] - item[0][0]) * (item[0][3] - item[0][1])) + return kind + + +def _draw_region( + response: Image.Image, + box: tuple[int, int, int, int], + color: tuple[int, int, int], + label: str = "", + font_family: str | None = None, +) -> None: + """Outline a region of the page and, if it is roomy enough, name it.""" + xmin, ymin, xmax, ymax = box + d = ImageDraw.Draw(response) + d.rectangle([(xmin, ymin), (xmax, ymax)], outline=color, width=2) + if label: + size = _fit_font_size(label, max(xmax - xmin, 1), max((ymax - ymin) // 4, 1), font_family, 8, 20) + font = _cached_font(font_family, size) + _draw_word(d, ((xmin + xmax) // 2, (ymin + ymax) // 2), label, font, color, anchor="mm") + + def _entry_geometry( entry: dict[str, Any], w: int, h: int ) -> tuple[list[tuple[float, float]], float, tuple[int, int, int, int]] | None: @@ -511,6 +571,7 @@ def _render( max_font_size: int, background_color: tuple[int, int, int], text_color: tuple[int, int, int], + draw_placeholders: bool, ) -> np.ndarray: """Render entries onto a blank page at one harmonized set of font sizes.""" h, w = _page_size(page) @@ -534,13 +595,60 @@ def _render( continue prepared.append((entry, polygon, angle, box, words, size)) - sizes = _harmonize([size for *_, size in prepared]) if prepared else [] + # A table draws its own grid, and a region holding no text of its own - a picture, a formula, + # a region the predictor was told to ignore - would otherwise leave a hole in the page. Both go + # down before the text, in a tone leaning towards the page itself so they stay quiet under it. + faint: tuple[int, int, int] = tuple((2 * back + fore) // 3 for back, fore in zip(background_color, text_color)) # type: ignore[assignment] + filled: list[tuple[int, int]] = [] + for _, _, _, box, entry_words, _ in prepared: + filled.extend([(word.x, word.y) for word in entry_words] or [((box[0] + box[2]) // 2, (box[1] + box[3]) // 2)]) + for table in page.get("tables") or []: + if not isinstance(table, dict): + continue + for region in [table, *(table.get("cells") or [])]: + geometry = _entry_geometry(region, w, h) if isinstance(region, dict) else None + if geometry is not None: + _draw_region(response, geometry[2], faint) + regions = [] + for region in page.get("layout") or []: + geometry = _entry_geometry(region, w, h) if isinstance(region, dict) else None + if geometry is None: + continue + regions.append((geometry[2], str(region.get("type", "")).strip())) + xmin, ymin, xmax, ymax = geometry[2] + if draw_placeholders and not any(xmin <= x <= xmax and ymin <= y <= ymax for x, y in filled): + _draw_region(response, geometry[2], faint, str(region.get("type", "")).strip(), font_family) + + sizes = [size for *_, size in prepared] + # A table cell is a structural box, not an ink box: it is as tall as its row, padding and all, + # so letting its text fill it renders a table twice the size of the page it sits on. The text + # a cell holds came off the same page, so the size the rest of the page uses is the better bet. + body = [size for (entry, *_), size in zip(prepared, sizes) if "row_start" not in entry] + if body: + sizes = [ + min(size, int(np.median(body))) if "row_start" in entry else size + for (entry, *_), size in zip(prepared, sizes) + ] + # Harmonizing inside a region rather than across the whole page is what keeps a section header + # its own size: it is often only a little larger than the body text, and a page-wide median + # would swallow it. Entries outside every region keep harmonizing with each other as before. + grouped: dict[str, list[int]] = {} + for index, (_, _, _, box, _, _) in enumerate(prepared): + centre = ((box[0] + box[2]) // 2, (box[1] + box[3]) // 2) + grouped.setdefault(_region_type(regions, centre), []).append(index) + for indices in grouped.values(): + for index, size in zip(indices, _harmonize([sizes[index] for index in indices])): + sizes[index] = size + for (entry, polygon, angle, box, words, _), size in zip(prepared, sizes): + centre = ((box[0] + box[2]) // 2, (box[1] + box[3]) // 2) + kind = _region_kind(_region_type(regions, centre)) + bold = kind in ("title", "sectionheader") try: if words: - _synthesize_line(response, words, size, font_family, min_font_size, text_color) + _synthesize_line(response, words, size, font_family, min_font_size, text_color, bold) else: - _synthesize_value(response, entry, polygon, angle, w, h, size, font_family, text_color) + _synthesize_value(response, entry, polygon, angle, w, h, size, font_family, text_color, bold) except Exception as exc: logging.warning(f"Could not render entry: {exc}") if draw_proba: @@ -557,6 +665,7 @@ def synthesize_page( max_font_size: int = 50, background_color: tuple[int, int, int] = (255, 255, 255), text_color: tuple[int, int, int] = (0, 0, 0), + draw_placeholders: bool = False, ) -> np.ndarray: """Draw the content of the element page (OCR response) on a blank page. @@ -568,18 +677,38 @@ def synthesize_page( max_font_size: maximum font size background_color: RGB color of the page background text_color: RGB color of the rendered text + draw_placeholders: if True, outline the layout regions holding no text of their own Returns: the synthesized page """ - lines = [ + entries = [ line for block in page.get("blocks") or [] if isinstance(block, dict) for line in block.get("lines") or [] if isinstance(line, dict) ] - return _render(page, lines, draw_proba, font_family, min_font_size, max_font_size, background_color, text_color) + # Words falling in a table are regrouped into its cells and removed from the blocks, so the + # cells are the only place that text still exists + entries += [ + cell + for table in page.get("tables") or [] + if isinstance(table, dict) + for cell in table.get("cells") or [] + if isinstance(cell, dict) + ] + return _render( + page, + entries, + draw_proba, + font_family, + min_font_size, + max_font_size, + background_color, + text_color, + draw_placeholders, + ) def synthesize_kie_page( @@ -590,6 +719,7 @@ def synthesize_kie_page( max_font_size: int = 50, background_color: tuple[int, int, int] = (255, 255, 255), text_color: tuple[int, int, int] = (0, 0, 0), + draw_placeholders: bool = False, ) -> np.ndarray: """Draw the content of the element page (KIE OCR response) on a blank page. @@ -601,6 +731,7 @@ def synthesize_kie_page( max_font_size: maximum font size background_color: RGB color of the page background text_color: RGB color of the rendered text + draw_placeholders: if True, outline the layout regions holding no text of their own Returns: the synthesized page @@ -612,5 +743,13 @@ def synthesize_kie_page( if isinstance(prediction, dict) ] return _render( - page, predictions, draw_proba, font_family, min_font_size, max_font_size, background_color, text_color + page, + predictions, + draw_proba, + font_family, + min_font_size, + max_font_size, + background_color, + text_color, + draw_placeholders, ) diff --git a/tests/common/test_utils_reconstitution.py b/tests/common/test_utils_reconstitution.py index 2afb52dca5..476b54a057 100644 --- a/tests/common/test_utils_reconstitution.py +++ b/tests/common/test_utils_reconstitution.py @@ -227,7 +227,7 @@ def test_synthesize_page_words_stay_apart_when_boxes_touch(): _assert_valid_render(render, (300, 900)) # Blank columns split the line into the gaps between letters and the two wider ones between - # the three words - if the words run together the widest gaps are just letter gaps + # the three words; if the words run together the widest gaps are just letter gaps columns = np.nonzero((render.min(axis=2) < 128).any(axis=0))[0] blanks = sorted((b - a - 1 for a, b in zip(columns, columns[1:]) if b - a > 1), reverse=True) assert blanks[1] > 1.5 * float(np.median(blanks)) @@ -290,3 +290,172 @@ def test_synthesize_page_line_holding_two_rows(): # baseline empties one row and fills the gap assert all((render[top : top + 30].min(axis=2) < 128).sum() > 50 for top in (120, 180)) assert (render[154:176] == 255).all() + + +def test_synthesize_page_table_cells(): + h_px, w_px = 1000, 1600 + words, x = [], 200.0 + for token in ("The", "table", "below", "lists", "the", "items"): + box_w = 12 * len(token) + words.append({ + "value": token, + "confidence": 1.0, + "geometry": ((x / w_px, 300 / h_px), ((x + box_w) / w_px, 322 / h_px)), + }) + x += box_w + 9 + cells = [] + for row, values in enumerate((("Product", "Price"), ("Widget", "9.99"))): + for col, value in enumerate(values): + x0, y0 = 200 + col * 400, 500 + row * 60 + cells.append({ + "value": value, + "confidence": 1.0, + "geometry": ((x0 / w_px, y0 / h_px), ((x0 + 380) / w_px, (y0 + 50) / h_px)), + "row_start": row, + "row_end": row, + "col_start": col, + "col_end": col, + }) + page = { + "dimensions": (h_px, w_px), + "blocks": [{"geometry": ((0, 0), (1, 1)), "lines": [{"geometry": ((0, 0), (1, 1)), "words": words}]}], + "tables": [ + { + "geometry": ((200 / w_px, 500 / h_px), (980 / w_px, 610 / h_px)), + "num_rows": 2, + "num_cols": 2, + "confidence": 0.9, + "cells": cells, + } + ], + } + render = reconstitution.synthesize_page(page) + _assert_valid_render(render, (h_px, w_px)) + + def ink_rows(y0, y1, x0, x1): + rows = np.nonzero((render[y0:y1, x0:x1].min(axis=2) < 128).any(axis=1))[0] + return rows.max() - rows.min() + 1 if len(rows) else 0 + + # every cell holds its text, and the grid is drawn in a lighter tone than the text + assert all(ink_rows(500, 610, x, x + 380) > 0 for x in (200, 600)) + assert ((render.min(axis=2) > 150) & (render.min(axis=2) < 220)).any() + # a table sized off its rows instead of off the page would be far taller than the body text + assert ink_rows(500, 548, 200, 980) <= 1.5 * ink_rows(295, 330, 200, 980) + + +def test_synthesize_page_layout_regions(): + page = { + "dimensions": (300, 400), + "blocks": [ + { + "geometry": ((0.1, 0.1), (0.4, 0.2)), + "lines": [ + { + "geometry": ((0.1, 0.1), (0.4, 0.2)), + "words": [{"value": "hello", "confidence": 0.9, "geometry": ((0.1, 0.1), (0.4, 0.2))}], + } + ], + } + ], + } + plain = reconstitution.synthesize_page(page) + _assert_valid_render(plain, (300, 400)) + + # a region wrapped around text that is already on the page must not add anything + covered = reconstitution.synthesize_page({ + **page, + "layout": [{"geometry": ((0.05, 0.05), (0.45, 0.25)), "type": "Text", "confidence": 0.9}], + }) + assert np.array_equal(plain, covered) + + # an empty one is only marked when the caller asks for it + empty = {**page, "layout": [{"geometry": ((0.5, 0.5), (0.95, 0.9)), "type": "Picture", "confidence": 0.8}]} + assert np.array_equal(plain, reconstitution.synthesize_page(empty)) + assert (reconstitution.synthesize_page(empty, draw_placeholders=True)[150:270, 200:380] < 255).any() + + +def test_synthesize_page_section_header_keeps_its_size(): + h_px, w_px = 1200, 1600 + + def line(text, y, size): + words, x = [], 200.0 + for token in text.split(): + box_w = int(size * 0.58 * len(token)) + words.append({ + "value": token, + "confidence": 1.0, + "geometry": ((x / w_px, y / h_px), ((x + box_w) / w_px, (y + size) / h_px)), + }) + x += box_w + int(size * 0.4) + xs = [c for word in words for c in (word["geometry"][0][0], word["geometry"][1][0])] + ys = [c for word in words for c in (word["geometry"][0][1], word["geometry"][1][1])] + return {"geometry": ((min(xs), min(ys)), (max(xs), max(ys))), "words": words} + + page = { + "dimensions": (h_px, w_px), + "blocks": [ + { + "geometry": ((0, 0), (1, 1)), + "lines": [ + line("Results", 300, 24), # a header only a little larger than the body + line("Results", 380, 22), + line("and the mean of each run is here", 420, 22), + line("with the deviation in brackets", 460, 22), + ], + } + ], + "layout": [ + {"geometry": ((0.11, 0.24), (0.40, 0.28)), "type": "Section-header", "confidence": 0.95}, + {"geometry": ((0.11, 0.30), (0.80, 0.42)), "type": "Text", "confidence": 0.96}, + ], + } + render = reconstitution.synthesize_page(page) + _assert_valid_render(render, (h_px, w_px)) + + def ink_width(y0, y1): + columns = np.nonzero((render[y0:y1, 190:400].min(axis=2) < 128).any(axis=0))[0] + return columns.max() - columns.min() + 1 if len(columns) else 0 + + # The same word sits in both, so its width follows the font size: harmonizing across the whole + # page instead of inside each region would render the header at the size of the body text + assert ink_width(290, 345) > 1.05 * ink_width(370, 412) + + +def test_synthesize_page_headings_are_bold(): + h_px, w_px = 600, 1200 + + def line(text, y, size=26): + words, x = [], 200.0 + for token in text.split(): + box_w = int(size * 0.58 * len(token)) + words.append({ + "value": token, + "confidence": 1.0, + "geometry": ((x / w_px, y / h_px), ((x + box_w) / w_px, (y + size) / h_px)), + }) + x += box_w + int(size * 0.4) + xs = [c for word in words for c in (word["geometry"][0][0], word["geometry"][1][0])] + ys = [c for word in words for c in (word["geometry"][0][1], word["geometry"][1][1])] + return {"geometry": ((min(xs), min(ys)), (max(xs), max(ys))), "words": words} + + page = { + "dimensions": (h_px, w_px), + "blocks": [{"geometry": ((0, 0), (1, 1)), "lines": [line("Heading", 200), line("Heading", 350)]}], + "layout": [ + {"geometry": ((0.12, 0.31), (0.50, 0.40)), "type": "Title", "confidence": 0.95}, + {"geometry": ((0.12, 0.56), (0.50, 0.66)), "type": "Text", "confidence": 0.95}, + ], + } + render = reconstitution.synthesize_page(page) + _assert_valid_render(render, (h_px, w_px)) + + def ink(y0, y1): + marked = render[y0:y1, 190:500].min(axis=2) < 128 + columns = np.nonzero(marked.any(axis=0))[0] + return marked.sum(), (columns.max() - columns.min() + 1 if len(columns) else 0) + + title, body = ink(190, 250), ink(340, 400) + # the same word in the same size of box: the title carries more ink without growing wider, + # which is weight rather than size + assert title[0] > 1.3 * body[0] + assert title[1] < 1.15 * body[1]