79 lines
2.8 KiB
Python
79 lines
2.8 KiB
Python
import logging
|
|
from typing import List
|
|
|
|
from fastapi import APIRouter, Depends, UploadFile, File, HTTPException, status
|
|
from google.api_core import exceptions as google_exceptions
|
|
|
|
from app.auth import current_active_user
|
|
from app.models import User as UserModel
|
|
from app.schemas.ocr import OcrExtractResponse
|
|
from app.core.gemini import extract_items_from_image_gemini, gemini_initialization_error, GeminiOCRService
|
|
from app.core.exceptions import (
|
|
OCRServiceUnavailableError,
|
|
OCRServiceConfigError,
|
|
OCRUnexpectedError,
|
|
OCRQuotaExceededError,
|
|
InvalidFileTypeError,
|
|
FileTooLargeError,
|
|
OCRProcessingError
|
|
)
|
|
from app.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
router = APIRouter()
|
|
ocr_service = GeminiOCRService()
|
|
|
|
@router.post(
|
|
"/extract-items",
|
|
response_model=OcrExtractResponse,
|
|
summary="Extract List Items via OCR (Gemini)",
|
|
tags=["OCR"]
|
|
)
|
|
async def ocr_extract_items(
|
|
current_user: UserModel = Depends(current_active_user),
|
|
image_file: UploadFile = File(..., description="Image file (JPEG, PNG, WEBP) of the shopping list or receipt."),
|
|
):
|
|
"""
|
|
Accepts an image upload, sends it to Gemini Flash with a prompt
|
|
to extract shopping list items, and returns the parsed items.
|
|
"""
|
|
# Check if Gemini client initialized correctly
|
|
if gemini_initialization_error:
|
|
logger.error("OCR endpoint called but Gemini client failed to initialize.")
|
|
raise OCRServiceUnavailableError(gemini_initialization_error)
|
|
|
|
logger.info(f"User {current_user.email} uploading image '{image_file.filename}' for OCR extraction.")
|
|
|
|
# --- File Validation ---
|
|
if image_file.content_type not in settings.ALLOWED_IMAGE_TYPES:
|
|
logger.warning(f"Invalid file type uploaded by {current_user.email}: {image_file.content_type}")
|
|
raise InvalidFileTypeError()
|
|
|
|
# Simple size check
|
|
contents = await image_file.read()
|
|
if len(contents) > settings.MAX_FILE_SIZE_MB * 1024 * 1024:
|
|
logger.warning(f"File too large uploaded by {current_user.email}: {len(contents)} bytes")
|
|
raise FileTooLargeError()
|
|
|
|
try:
|
|
# Call the Gemini helper function
|
|
extracted_items = await extract_items_from_image_gemini(
|
|
image_bytes=contents,
|
|
mime_type=image_file.content_type
|
|
)
|
|
|
|
logger.info(f"Successfully extracted {len(extracted_items)} items for user {current_user.email}.")
|
|
return OcrExtractResponse(extracted_items=extracted_items)
|
|
|
|
except OCRServiceUnavailableError:
|
|
raise OCRServiceUnavailableError()
|
|
except OCRServiceConfigError:
|
|
raise OCRServiceConfigError()
|
|
except OCRQuotaExceededError:
|
|
raise OCRQuotaExceededError()
|
|
except Exception as e:
|
|
raise OCRProcessingError(str(e))
|
|
|
|
finally:
|
|
# Ensure file handle is closed
|
|
await image_file.close() |