| |
| """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() |
|
|
