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