Spaces:
Build error
Build error
| from __future__ import annotations | |
| import importlib.util | |
| import math | |
| from dataclasses import dataclass | |
| import numpy as np | |
| from .models import Device, Point | |
| class OcrUnavailable(RuntimeError): | |
| pass | |
| class OcrLine: | |
| text: str | |
| confidence: float | |
| quad: list[Point] | |
| def is_ocr_available() -> bool: | |
| return importlib.util.find_spec("rapidocr") is not None and ( | |
| importlib.util.find_spec("onnxruntime") is not None | |
| ) | |
| class RapidOcrEngine: | |
| def __init__(self, device: Device = Device.AUTO) -> None: | |
| if not is_ocr_available(): | |
| raise OcrUnavailable( | |
| "ONNX Runtime is missing. Install the cpu or cuda project extra." | |
| ) | |
| from rapidocr import RapidOCR | |
| # RapidOCR 3.9 defaults to PP-OCRv6-small detection and recognition. | |
| self._engine = RapidOCR(params={"Global.use_cls": False}) | |
| def recognize(self, rgba: np.ndarray, confidence: float) -> list[OcrLine]: | |
| rgb = rgba[:, :, :3] | |
| result = self._engine(rgb) | |
| boxes = getattr(result, "boxes", None) | |
| txts = getattr(result, "txts", None) | |
| scores = getattr(result, "scores", None) | |
| if boxes is None or txts is None or scores is None: | |
| return [] | |
| lines: list[OcrLine] = [] | |
| for box, text, score in zip(boxes, txts, scores, strict=True): | |
| value = float(score) | |
| clean = str(text).strip() | |
| if value < confidence or not clean: | |
| continue | |
| quad = [Point(x=max(0, float(x)), y=max(0, float(y))) for x, y in box] | |
| if len(quad) != 4 or polygon_area(quad) < 4: | |
| continue | |
| lines.append(OcrLine(text=clean, confidence=value, quad=quad)) | |
| return sorted(lines, key=lambda line: (quad_box(line.quad)[1], quad_box(line.quad)[0])) | |
| def polygon_area(points: list[Point]) -> float: | |
| return abs( | |
| sum( | |
| point.x * points[(index + 1) % len(points)].y | |
| - points[(index + 1) % len(points)].x * point.y | |
| for index, point in enumerate(points) | |
| ) | |
| / 2 | |
| ) | |
| def quad_box(quad: list[Point]) -> tuple[float, float, float, float]: | |
| xs = [point.x for point in quad] | |
| ys = [point.y for point in quad] | |
| return min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys) | |
| def quad_rotation(quad: list[Point]) -> float: | |
| first, second = quad[0], quad[1] | |
| return math.degrees(math.atan2(second.y - first.y, second.x - first.x)) | |
| def snap_rotation(angle: float, targets: list[float], tolerance: float) -> float: | |
| nearest = min(targets, key=lambda target: angular_distance(angle, target)) | |
| if angular_distance(angle, nearest) <= tolerance: | |
| return float(nearest) | |
| return angle | |
| def angular_distance(first: float, second: float) -> float: | |
| return abs((first - second + 180) % 360 - 180) | |