app/services/ocr.py

"""Обёртка над fast-plate-ocr с устойчивостью к разным версиям API."""
import tempfile
from pathlib import Path

import cv2
import numpy as np
from fast_plate_ocr import LicensePlateRecognizer

# Доступные модели от лучшей к самой быстрой:
#   "global-plates-mobile-vit-v2-model" — самая точная, ~30 ms на CPU
#   "cct-s-v2-global-model"             — баланс, ~10 ms
#   "cct-xs-v2-global-model"            — самая быстрая, ~5 ms
DEFAULT_OCR_MODEL = "global-plates-mobile-vit-v2-model"


class PlateOCR:
    def __init__(self, model_name: str = DEFAULT_OCR_MODEL):
        self.model = LicensePlateRecognizer(model_name)

    def _extract_text(self, pred) -> str:
        """API менялся между версиями fast-plate-ocr."""
        if isinstance(pred, str):
            return pred
        for attr in ("plate", "text", "label", "plate_text"):
            v = getattr(pred, attr, None)
            if isinstance(v, str):
                return v
        return str(pred)

    def recognize(self, image: np.ndarray) -> str:
        """Принимает np.ndarray (BGR), возвращает raw текст."""
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
            tmp_path = Path(f.name)
        try:
            cv2.imwrite(str(tmp_path), image)
            preds = self.model.run(str(tmp_path))
            if not preds:
                return ""
            return self._extract_text(preds[0])
        finally:
            tmp_path.unlink(missing_ok=True)