ocr_studio/tiling.py

from __future__ import annotations

from collections.abc import Iterator

from PIL import Image

from ocr_studio.config import TILE_HEIGHT, TILE_OVERLAP
from ocr_studio.spotting import TextSpan, boxes_iou


def should_tile(image: Image.Image, enabled: bool) -> bool:
    if not enabled:
        return False
    return image.height > TILE_HEIGHT + 80 or image.width > 2200


def iter_tiles(
    image: Image.Image,
    tile_height: int = TILE_HEIGHT,
    overlap: int = TILE_OVERLAP,
) -> Iterator[tuple[Image.Image, int, int]]:
    width, height = image.size
    if height <= tile_height + 40:
        yield image, 0, 0
        return

    step = max(240, tile_height - overlap)
    y = 0
    while y < height:
        y1 = min(height, y + tile_height)
        yield image.crop((0, y, width, y1)), 0, y
        if y1 >= height:
            break
        y += step


def offset_spans(spans: list[TextSpan], origin_x: int, origin_y: int) -> list[TextSpan]:
    shifted: list[TextSpan] = []
    for span in spans:
        if span.box is None:
            shifted.append(span)
            continue
        x0, y0, x1, y1 = span.box
        shifted.append(
            TextSpan(
                text=span.text,
                box=(x0 + origin_x, y0 + origin_y, x1 + origin_x, y1 + origin_y),
            )
        )
    return shifted


def merge_tile_spans(spans: list[TextSpan]) -> list[TextSpan]:
    kept: list[TextSpan] = []
    for span in sorted(spans, key=lambda item: (item.box or (0, 0, 0, 0))[1]):
        duplicate = False
        for existing in kept:
            if not span.box or not existing.box:
                continue
            if boxes_iou(span.box, existing.box) < 0.45:
                continue
            if span.text.strip() == existing.text.strip() or span.text.strip() in existing.text:
                duplicate = True
                break
        if not duplicate:
            kept.append(span)
    return kept


def join_tile_text(chunks: list[str]) -> str:
    cleaned = [chunk.strip() for chunk in chunks if chunk and chunk.strip()]
    if not cleaned:
        return ""
    merged = [cleaned[0]]
    for chunk in cleaned[1:]:
        previous = merged[-1]
        overlap = _overlap_suffix(previous, chunk)
        if overlap:
            merged[-1] = previous + chunk[overlap:]
        else:
            merged.append(chunk)
    return "\n".join(merged)


def _overlap_suffix(left: str, right: str, max_len: int = 80) -> int:
    limit = min(len(left), len(right), max_len)
    for size in range(limit, 12, -1):
        if left[-size:] == right[:size]:
            return size
    return 0