mirror of
https://github.com/serty2005/rmser.git
synced 2026-02-04 19:02:33 -06:00
102 lines
3.7 KiB
Python
102 lines
3.7 KiB
Python
import logging
|
||
from typing import List
|
||
|
||
from fastapi import FastAPI, File, UploadFile, HTTPException
|
||
from pydantic import BaseModel
|
||
import cv2
|
||
import numpy as np
|
||
|
||
# Импортируем модули
|
||
from imgproc import preprocess_image
|
||
from parser import parse_receipt_text, ParsedItem
|
||
from ocr import ocr_engine
|
||
# Импортируем новый модуль
|
||
from qr_manager import detect_and_decode_qr, fetch_data_from_api
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
app = FastAPI(title="RMSER OCR Service (Hybrid: QR + OCR)")
|
||
|
||
class RecognitionResult(BaseModel):
|
||
source: str # 'qr_api' или 'ocr'
|
||
items: List[ParsedItem]
|
||
raw_text: str = ""
|
||
|
||
@app.get("/health")
|
||
def health_check():
|
||
return {"status": "ok"}
|
||
|
||
@app.post("/recognize", response_model=RecognitionResult)
|
||
async def recognize_receipt(image: UploadFile = File(...)):
|
||
"""
|
||
1. Попытка найти QR-код.
|
||
2. Если QR найден -> запрос к API -> возврат идеальных данных.
|
||
3. Если QR не найден -> Preprocessing -> OCR -> Regex Parsing.
|
||
"""
|
||
logger.info(f"Received file: {image.filename}, content_type: {image.content_type}")
|
||
|
||
if not image.content_type.startswith("image/"):
|
||
raise HTTPException(status_code=400, detail="File must be an image")
|
||
|
||
try:
|
||
# Читаем байты
|
||
content = await image.read()
|
||
|
||
# Конвертируем в numpy для работы (нужен и для QR, и для OCR)
|
||
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("Attempting 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"Successfully retrieved {len(api_items)} items via API.")
|
||
return RecognitionResult(
|
||
source="qr_api",
|
||
items=api_items,
|
||
raw_text=f"QR Content: {qr_raw}"
|
||
)
|
||
else:
|
||
logger.warning("QR found but API failed to return items. Falling back to OCR.")
|
||
else:
|
||
logger.info("QR code not found. Falling back to OCR.")
|
||
|
||
# --- ЭТАП 2: OCR Strategy (Fallback) ---
|
||
|
||
# 1. Image Processing (получаем бинарное изображение)
|
||
# Передаем исходные байты, так как функция внутри декодирует их заново
|
||
# (можно оптимизировать, но оставим совместимость с текущим кодом)
|
||
processed_img = preprocess_image(content)
|
||
|
||
# 2. OCR
|
||
full_text = ocr_engine.recognize(processed_img)
|
||
|
||
# 3. Parsing
|
||
ocr_items = parse_receipt_text(full_text)
|
||
|
||
return RecognitionResult(
|
||
source="ocr",
|
||
items=ocr_items,
|
||
raw_text=full_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) |