247 lines
7.9 KiB
Python
247 lines
7.9 KiB
Python
"""OCR service: file upload, result retrieval, list, apply-to-vehicle, and async processing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
from sqlalchemy import and_, func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import settings
|
|
from app.models.ocr_result import OCRResult, OCRStatus
|
|
from app.models.vehicle import Vehicle
|
|
from app.utils.openrouter import perform_ocr
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Confidence threshold: below this, status is set to manual_review
|
|
CONFIDENCE_THRESHOLD = 0.7
|
|
|
|
# Allowed MIME types for OCR uploads
|
|
ALLOWED_MIME_TYPES = {"image/png", "image/jpeg", "image/jpg", "image/webp", "image/gif"}
|
|
|
|
|
|
def validate_mime_type(mime_type: str) -> bool:
|
|
"""Check if the MIME type is allowed for OCR uploads."""
|
|
return mime_type in ALLOWED_MIME_TYPES
|
|
|
|
|
|
def validate_file_size(file_size_bytes: int) -> bool:
|
|
"""Check if the file size is within the configured limit."""
|
|
max_bytes = settings.MAX_FILE_SIZE_MB * 1024 * 1024
|
|
return file_size_bytes <= max_bytes
|
|
|
|
|
|
async def upload_file(
|
|
db: AsyncSession,
|
|
file_bytes: bytes,
|
|
file_name: str,
|
|
mime_type: str,
|
|
vehicle_id: uuid.UUID | None = None,
|
|
) -> OCRResult:
|
|
"""Save uploaded file to disk and create an OCRResult record with status=pending.
|
|
|
|
Validates MIME type and file size before saving.
|
|
"""
|
|
if not validate_mime_type(mime_type):
|
|
raise ValueError(f"Invalid MIME type: {mime_type}. Allowed: {ALLOWED_MIME_TYPES}")
|
|
|
|
if not validate_file_size(len(file_bytes)):
|
|
raise ValueError(
|
|
f"File size exceeds limit of {settings.MAX_FILE_SIZE_MB} MB"
|
|
)
|
|
|
|
# Ensure upload directory exists
|
|
upload_dir = settings.UPLOAD_DIR
|
|
os.makedirs(upload_dir, exist_ok=True)
|
|
|
|
# Generate unique filename
|
|
file_ext = os.path.splitext(file_name)[1] or ".png"
|
|
unique_name = f"{uuid.uuid4().hex}{file_ext}"
|
|
file_path = os.path.join(upload_dir, unique_name)
|
|
|
|
# Write file to disk
|
|
with open(file_path, "wb") as f:
|
|
f.write(file_bytes)
|
|
|
|
# Create OCR result record
|
|
ocr_result = OCRResult(
|
|
vehicle_id=vehicle_id,
|
|
file_path=file_path,
|
|
file_name=file_name,
|
|
mime_type=mime_type,
|
|
status=OCRStatus.pending,
|
|
)
|
|
db.add(ocr_result)
|
|
await db.flush()
|
|
await db.refresh(ocr_result)
|
|
|
|
return ocr_result
|
|
|
|
|
|
async def get_result(db: AsyncSession, result_id: uuid.UUID) -> OCRResult | None:
|
|
"""Get a single OCR result by ID."""
|
|
stmt = select(OCRResult).where(OCRResult.id == result_id)
|
|
result = await db.execute(stmt)
|
|
return result.scalar_one_or_none()
|
|
|
|
|
|
async def list_results(
|
|
db: AsyncSession,
|
|
vehicle_id: uuid.UUID | None = None,
|
|
page: int = 1,
|
|
page_size: int = 20,
|
|
) -> tuple[list[OCRResult], int]:
|
|
"""List OCR results, optionally filtered by vehicle_id, with pagination."""
|
|
conditions = []
|
|
if vehicle_id is not None:
|
|
conditions.append(OCRResult.vehicle_id == vehicle_id)
|
|
|
|
# Count query
|
|
count_stmt = select(func.count(OCRResult.id))
|
|
if conditions:
|
|
count_stmt = count_stmt.where(and_(*conditions))
|
|
total_result = await db.execute(count_stmt)
|
|
total = total_result.scalar_one()
|
|
|
|
# Data query
|
|
data_stmt = select(OCRResult).order_by(OCRResult.created_at.desc())
|
|
if conditions:
|
|
data_stmt = data_stmt.where(and_(*conditions))
|
|
|
|
offset = (page - 1) * page_size
|
|
data_stmt = data_stmt.offset(offset).limit(page_size)
|
|
|
|
result = await db.execute(data_stmt)
|
|
items = list(result.scalars().all())
|
|
|
|
return items, total
|
|
|
|
|
|
async def apply_to_vehicle(
|
|
db: AsyncSession, result_id: uuid.UUID
|
|
) -> tuple[OCRResult, Vehicle, list[str]]:
|
|
"""Apply OCR structured data to the linked vehicle.
|
|
|
|
Maps OCR fields to vehicle fields:
|
|
brand → make, model → model, vin → fin,
|
|
first_registration → first_registration, mileage → mileage_km,
|
|
power_kw → power_kw, fuel_type → fuel_type
|
|
|
|
Returns (ocr_result, vehicle, updated_fields).
|
|
Raises ValueError if OCR result not found, no vehicle linked, or no structured data.
|
|
"""
|
|
ocr_result = await get_result(db, result_id)
|
|
if ocr_result is None:
|
|
raise ValueError("OCR result not found")
|
|
|
|
if ocr_result.vehicle_id is None:
|
|
raise ValueError("No vehicle linked to this OCR result")
|
|
|
|
if not ocr_result.structured_data:
|
|
raise ValueError("No structured data available to apply")
|
|
|
|
# Fetch vehicle
|
|
stmt = select(Vehicle).where(
|
|
and_(Vehicle.id == ocr_result.vehicle_id, Vehicle.deleted_at.is_(None))
|
|
)
|
|
vehicle_result = await db.execute(stmt)
|
|
vehicle = vehicle_result.scalar_one_or_none()
|
|
if vehicle is None:
|
|
raise ValueError("Linked vehicle not found")
|
|
|
|
data = ocr_result.structured_data
|
|
updated_fields: list[str] = []
|
|
|
|
# Map OCR fields to vehicle fields
|
|
field_mapping = {
|
|
"brand": "make",
|
|
"model": "model",
|
|
"vin": "fin",
|
|
"first_registration": "first_registration",
|
|
"mileage": "mileage_km",
|
|
"power_kw": "power_kw",
|
|
"fuel_type": "fuel_type",
|
|
}
|
|
|
|
for ocr_field, vehicle_field in field_mapping.items():
|
|
value = data.get(ocr_field)
|
|
if value is not None and value != "":
|
|
# Parse first_registration to date
|
|
if ocr_field == "first_registration" and isinstance(value, str):
|
|
try:
|
|
from datetime import datetime as dt
|
|
parsed = dt.strptime(value, "%d.%m.%Y").date()
|
|
setattr(vehicle, vehicle_field, parsed)
|
|
updated_fields.append(vehicle_field)
|
|
continue
|
|
except ValueError:
|
|
try:
|
|
from datetime import date
|
|
parsed = date.fromisoformat(value)
|
|
setattr(vehicle, vehicle_field, parsed)
|
|
updated_fields.append(vehicle_field)
|
|
continue
|
|
except ValueError:
|
|
logger.warning("Could not parse first_registration: %s", value)
|
|
continue
|
|
|
|
setattr(vehicle, vehicle_field, value)
|
|
updated_fields.append(vehicle_field)
|
|
|
|
await db.flush()
|
|
await db.refresh(vehicle)
|
|
|
|
return ocr_result, vehicle, updated_fields
|
|
|
|
|
|
async def process_ocr(db: AsyncSession, result_id: uuid.UUID) -> OCRResult:
|
|
"""Process an OCR result: read file, call OpenRouter, update result.
|
|
|
|
This is the async processing function called by the background task.
|
|
Sets status to 'completed' if confidence >= threshold, else 'manual_review'.
|
|
Sets status to 'failed' on error.
|
|
"""
|
|
ocr_result = await get_result(db, result_id)
|
|
if ocr_result is None:
|
|
raise ValueError(f"OCR result {result_id} not found")
|
|
|
|
# Update status to processing
|
|
ocr_result.status = OCRStatus.processing
|
|
await db.flush()
|
|
|
|
try:
|
|
# Read file from disk
|
|
with open(ocr_result.file_path, "rb") as f:
|
|
image_bytes = f.read()
|
|
|
|
# Call OpenRouter vision model
|
|
ocr_output = await perform_ocr(
|
|
image_bytes=image_bytes,
|
|
mime_type=ocr_result.mime_type,
|
|
)
|
|
|
|
# Update OCR result with extracted data
|
|
ocr_result.raw_text = ocr_output.get("raw_text", "")
|
|
ocr_result.structured_data = ocr_output.get("structured_data", {})
|
|
ocr_result.confidence_score = ocr_output.get("confidence_score", 0.0)
|
|
|
|
# Set status based on confidence threshold
|
|
if ocr_result.confidence_score >= CONFIDENCE_THRESHOLD:
|
|
ocr_result.status = OCRStatus.completed
|
|
else:
|
|
ocr_result.status = OCRStatus.manual_review
|
|
|
|
except Exception as exc:
|
|
logger.error("OCR processing failed for %s: %s", result_id, exc)
|
|
ocr_result.status = OCRStatus.failed
|
|
ocr_result.error_message = str(exc)
|
|
|
|
await db.flush()
|
|
await db.refresh(ocr_result)
|
|
return ocr_result
|