rmser/ocr-service/app/main.py

import logging
import os
from typing import List

from fastapi import FastAPI, File, UploadFile, HTTPException
from pydantic import BaseModel
import cv2
import numpy as np

# Импортируем модули
from app.schemas.models import ParsedItem, RecognitionResult
from app.services.qr import detect_and_decode_qr, fetch_data_from_api, extract_fiscal_data
# Импортируем новый модуль
from app.services.ocr import yandex_engine
from app.services.llm import llm_parser
from app.services.excel import extract_text_from_excel

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

app = FastAPI(title="RMSER OCR Service (Cloud-only: QR + Yandex + GigaChat)")

@app.get("/health")
def health_check():
    return {"status": "ok"}

@app.post("/recognize", response_model=RecognitionResult)
async def recognize_receipt(image: UploadFile = File(...)):
    """
    Стратегия:
    1. Excel файл (.xlsx) -> Извлечение текста -> LLM парсинг
    2. QR Code + FNS API (Приоритет 1 - Идеальная точность)
    3. Yandex Vision OCR + LLM (Приоритет 2 - Высокая точность, если настроен)
    Если ничего не найдено, возвращает пустой результат.
    """
    logger.info(f"Received file: {image.filename}, content_type: {image.content_type}")

    # Проверка на Excel файл
    if image.filename and image.filename.lower().endswith('.xlsx'):
        logger.info("Processing Excel file...")
        try:
            content = await image.read()
            excel_text = extract_text_from_excel(content)
            
            if excel_text:
                logger.info(f"Excel text extracted, length: {len(excel_text)}")
                logger.info("Calling LLM Manager to parse Excel text...")
                excel_result = llm_parser.parse_receipt(excel_text)
                
                return RecognitionResult(
                    source="excel_llm",
                    items=excel_result["items"],
                    raw_text=excel_text,
                    doc_number=excel_result["doc_number"],
                    doc_date=excel_result["doc_date"]
                )
            else:
                logger.warning("Excel file is empty or contains no text")
                return RecognitionResult(
                    source="none",
                    items=[],
                    raw_text=""
                )
        except Exception as e:
            logger.error(f"Error processing Excel file: {e}", exc_info=True)
            raise HTTPException(status_code=500, detail=f"Error processing Excel file: {str(e)}")

    # Проверка на изображение
    if not image.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image or .xlsx file")

    try:
        # Читаем сырые байты
        content = await image.read()
        
        # Конвертируем в numpy для QR
        nparr = np.frombuffer(content, np.uint8)
        original_cv_image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        if original_cv_image is None:
             raise HTTPException(status_code=400, detail="Invalid image data")

        # --- ЭТАП 1: QR Code Strategy ---
        logger.info("--- Stage 1: QR Code Detection ---")
        qr_raw = detect_and_decode_qr(original_cv_image)
        
        if qr_raw:
            logger.info("QR found! Fetching data from API...")
            api_items = fetch_data_from_api(qr_raw)
            
            if api_items:
                logger.info(f"Success: Retrieved {len(api_items)} items via QR API.")
                return RecognitionResult(
                    source="qr_api",
                    items=api_items,
                    raw_text=f"QR Content: {qr_raw}"
                )
            else:
                logger.warning("QR found but API failed. Falling back to OCR.")
        else:
            logger.info("QR code not found. Proceeding to OCR.")

        # --- ЭТАП 2: OCR + Virtual QR Strategy ---
        if yandex_engine.is_configured():
            logger.info("--- Stage 2: Yandex Vision OCR + Virtual QR ---")
            yandex_text = yandex_engine.recognize(content)

            if yandex_text and len(yandex_text) > 10:
                logger.info(f"OCR success. Raw text length: {len(yandex_text)}")

                # Попытка собрать виртуальный QR из текста
                virtual_qr = extract_fiscal_data(yandex_text)
                if virtual_qr:
                    logger.info(f"Virtual QR constructed: {virtual_qr}")
                    api_items = fetch_data_from_api(virtual_qr)
                    if api_items:
                        logger.info(f"Success: Retrieved {len(api_items)} items via Virtual QR API.")
                        return RecognitionResult(
                            source="virtual_qr_api",
                            items=api_items,
                            raw_text=yandex_text
                        )

                # Вызываем LLM для парсинга текста
                logger.info("Calling LLM Manager to parse text...")
                yandex_result = llm_parser.parse_receipt(yandex_text)

                return RecognitionResult(
                    source="yandex_vision_llm",
                    items=yandex_result["items"],
                    raw_text=yandex_text,
                    doc_number=yandex_result["doc_number"],
                    doc_date=yandex_result["doc_date"]
                )
            else:
                logger.warning("Yandex Vision returned empty text or failed. No fallback available.")
                return RecognitionResult(
                    source="none",
                    items=[],
                    raw_text=""
                )
        else:
            logger.info("Yandex Vision credentials not set. No OCR available.")
            return RecognitionResult(
                source="none",
                items=[],
                raw_text=""
            )

    except Exception as e:
        logger.error(f"Error processing request: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5000)