mirror of
https://github.com/serty2005/rmser.git
synced 2026-02-04 19:02:33 -06:00
114 lines
4.6 KiB
Python
114 lines
4.6 KiB
Python
from abc import ABC, abstractmethod
|
||
from typing import Optional
|
||
import os
|
||
import json
|
||
import base64
|
||
import logging
|
||
import requests
|
||
|
||
from app.services.auth import yandex_auth
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
VISION_URL = "https://ocr.api.cloud.yandex.net/ocr/v1/recognizeText"
|
||
|
||
class OCREngine(ABC):
|
||
@abstractmethod
|
||
def recognize(self, image_bytes: bytes) -> str:
|
||
"""Распознает текст из изображения и возвращает строку."""
|
||
pass
|
||
|
||
@abstractmethod
|
||
def is_configured(self) -> bool:
|
||
"""Проверяет, настроен ли движок (наличие ключей/настроек)."""
|
||
pass
|
||
|
||
class YandexOCREngine(OCREngine):
|
||
def __init__(self):
|
||
self.folder_id = os.getenv("YANDEX_FOLDER_ID")
|
||
|
||
if not yandex_auth.is_configured() or not self.folder_id:
|
||
logger.warning("Yandex OCR credentials (YANDEX_OAUTH_TOKEN, YANDEX_FOLDER_ID) not set. Yandex OCR will be unavailable.")
|
||
|
||
def is_configured(self) -> bool:
|
||
return yandex_auth.is_configured() and bool(self.folder_id)
|
||
|
||
def recognize(self, image_bytes: bytes) -> str:
|
||
"""
|
||
Отправляет изображение в Yandex Vision и возвращает полный текст.
|
||
"""
|
||
if not self.is_configured():
|
||
logger.error("Yandex credentials missing.")
|
||
return ""
|
||
|
||
iam_token = yandex_auth.get_iam_token()
|
||
if not iam_token:
|
||
return ""
|
||
|
||
# 1. Кодируем в Base64
|
||
b64_image = base64.b64encode(image_bytes).decode("utf-8")
|
||
|
||
# 2. Формируем тело запроса
|
||
# Используем модель 'page' (для документов) и '*' для автоопределения языка
|
||
payload = {
|
||
"mimeType": "JPEG", # Yandex переваривает и PNG под видом JPEG часто, но лучше быть аккуратным.
|
||
# В идеале определять mime-type из файла, но JPEG - безопасный дефолт для фото.
|
||
"languageCodes": ["*"],
|
||
"model": "page",
|
||
"content": b64_image
|
||
}
|
||
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {iam_token}",
|
||
"x-folder-id": self.folder_id,
|
||
"x-data-logging-enabled": "true"
|
||
}
|
||
|
||
# 3. Отправляем запрос
|
||
try:
|
||
logger.info("Sending request to Yandex Vision OCR...")
|
||
response = requests.post(VISION_URL, headers=headers, json=payload, timeout=20)
|
||
|
||
# Если 401 Unauthorized, возможно токен протух раньше времени (редко, но бывает)
|
||
if response.status_code == 401:
|
||
logger.warning("Got 401 from Yandex. Retrying with fresh token...")
|
||
yandex_auth.reset_token() # сброс кэша
|
||
iam_token = yandex_auth.get_iam_token()
|
||
if iam_token:
|
||
headers["Authorization"] = f"Bearer {iam_token}"
|
||
response = requests.post(VISION_URL, headers=headers, json=payload, timeout=20)
|
||
|
||
response.raise_for_status()
|
||
result_json = response.json()
|
||
|
||
# 4. Парсим ответ
|
||
# Структура: result -> textAnnotation -> fullText
|
||
# Или (если fullText нет) blocks -> lines -> text
|
||
|
||
text_annotation = result_json.get("result", {}).get("textAnnotation", {})
|
||
|
||
if not text_annotation:
|
||
logger.warning("Yandex returned success but no textAnnotation found.")
|
||
return ""
|
||
|
||
# Самый простой способ - взять fullText, он обычно склеен с \n
|
||
full_text = text_annotation.get("fullText", "")
|
||
|
||
if not full_text:
|
||
# Фолбэк: если fullText пуст, собираем вручную по блокам
|
||
logger.info("fullText empty, assembling from blocks...")
|
||
lines_text = []
|
||
for block in text_annotation.get("blocks", []):
|
||
for line in block.get("lines", []):
|
||
lines_text.append(line.get("text", ""))
|
||
full_text = "\n".join(lines_text)
|
||
|
||
return full_text
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error during Yandex Vision request: {e}")
|
||
return ""
|
||
|
||
# Глобальный инстанс
|
||
yandex_engine = YandexOCREngine() |