Spaces:
Build error
Build error
File size: 2,882 Bytes
3c2cd23 | 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 | 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
@dataclass(slots=True)
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)
|