Spaces:
Running on Zero
Running on Zero
File size: 2,716 Bytes
98e1d9f | 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 | 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
|