File size: 4,111 Bytes
6f683f1 | 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 | #!/usr/bin/env python3
"""Split long RTL OCR lines into CTC-safe Persian word-group crops."""
from __future__ import annotations
import argparse
from pathlib import Path
import cv2
import numpy as np
from PIL import Image
def zero_runs(mask: np.ndarray) -> list[tuple[int, int]]:
runs = []
start = None
for index, value in enumerate(mask.tolist() + [False]):
if value and start is None:
start = index
elif not value and start is not None:
runs.append((start, index))
start = None
return runs
def word_intervals(image: Image.Image, count: int) -> list[tuple[int, int]] | None:
if count == 1:
return [(0, image.width)]
gray = np.asarray(image.convert("L"))
_, ink = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
projection = (ink > 0).sum(axis=0)
active = np.flatnonzero(projection > max(1, image.height // 80))
if active.size == 0:
return None
left, right = int(active[0]), int(active[-1]) + 1
gaps = [
(a, b) for a, b in zero_runs(projection[left:right] <= max(1, image.height // 100))
if b - a >= 2 and a > 0 and b < right - left
]
if len(gaps) < count - 1:
return None
chosen = sorted(
((a + left, b + left) for a, b in sorted(gaps, key=lambda run: run[1] - run[0], reverse=True)[: count - 1]),
key=lambda run: run[0],
)
cuts = [left] + [(a + b) // 2 for a, b in chosen] + [right]
intervals = [(cuts[i], cuts[i + 1]) for i in range(len(cuts) - 1)]
return list(reversed(intervals))
def groups(words: list[str], intervals: list[tuple[int, int]], max_chars: int):
current_words: list[str] = []
current_boxes: list[tuple[int, int]] = []
for word, box in zip(words, intervals):
candidate = " ".join(current_words + [word])
if current_words and len(candidate) > max_chars:
yield " ".join(current_words), current_boxes
current_words, current_boxes = [], []
if len(word) > max_chars:
return
current_words.append(word)
current_boxes.append(box)
if current_words:
yield " ".join(current_words), current_boxes
def process_list(source: Path, destination: Path, image_dir: Path, root: Path, max_chars: int) -> tuple[int, int]:
output = []
rejected = 0
for line_index, raw in enumerate(source.read_text(encoding="utf-8").splitlines()):
relative, text = raw.split("\t", 1)
words = text.split()
with Image.open(root / relative) as opened:
image = opened.convert("RGB")
intervals = word_intervals(image, len(words))
if intervals is None:
rejected += 1
continue
for group_index, (label, boxes) in enumerate(groups(words, intervals, max_chars)):
x0 = max(0, min(x[0] for x in boxes) - 3)
x1 = min(image.width, max(x[1] for x in boxes) + 3)
crop = image.crop((x0, 0, x1, image.height))
path = image_dir / f"{line_index:07d}-{group_index:03d}.png"
crop.save(path)
output.append(f"{path.relative_to(root)}\t{label}")
destination.write_text("\n".join(output) + "\n", encoding="utf-8")
return len(output), rejected
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--data-root", type=Path, required=True)
parser.add_argument("--pilot-dir", type=Path, required=True)
parser.add_argument("--max-chars", type=int, default=25)
args = parser.parse_args()
for split in ("train", "val"):
image_dir = args.pilot_dir / "images" / f"segments-{split}"
image_dir.mkdir(parents=True, exist_ok=True)
kept, rejected = process_list(
args.pilot_dir / f"{split}_list.txt",
args.pilot_dir / f"{split}_segments.txt",
image_dir,
args.data_root,
args.max_chars,
)
print(f"{split}: segments={kept} rejected_lines={rejected}")
if __name__ == "__main__":
main()
|