Fix startup and Gemini provider handling
This commit is contained in:
@@ -7,6 +7,7 @@ backups/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.venv/
|
||||
.venv312/
|
||||
venv/
|
||||
*.egg-info/
|
||||
.claude/
|
||||
@@ -14,3 +15,4 @@ venv/
|
||||
*.jpeg
|
||||
*.jpg
|
||||
*.png
|
||||
tmp_samples/
|
||||
|
||||
@@ -12,6 +12,10 @@ from routers import documents, review, search, upload, auth
|
||||
from services.auth_service import get_session, cleanup_expired_sessions
|
||||
|
||||
|
||||
Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
|
||||
Path("data").mkdir(exist_ok=True)
|
||||
|
||||
|
||||
class CachedStaticFiles(StaticFiles):
|
||||
"""StaticFiles that sets a long Cache-Control header so browsers
|
||||
don't re-download the same image on every refresh."""
|
||||
@@ -29,8 +33,6 @@ class CachedStaticFiles(StaticFiles):
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True)
|
||||
Path("data").mkdir(exist_ok=True)
|
||||
create_tables()
|
||||
try:
|
||||
cleanup_expired_sessions()
|
||||
|
||||
@@ -5,6 +5,7 @@ google-genai>=1.70.0
|
||||
easyocr>=1.7.0
|
||||
pymupdf>=1.24.0
|
||||
jinja2>=3.1.4
|
||||
python-dotenv>=1.0.1
|
||||
python-multipart>=0.0.9
|
||||
aiofiles>=23.0.0
|
||||
pillow>=10.0.0
|
||||
|
||||
+20
-28
@@ -45,7 +45,7 @@ def require_admin(request: Request):
|
||||
|
||||
@router.get("/login", response_class=HTMLResponse)
|
||||
async def login_get(request: Request):
|
||||
return templates.TemplateResponse(request=request, name="login.html", context={"error": None})
|
||||
return templates.TemplateResponse("login.html", {"request": request, "error": None})
|
||||
|
||||
@router.post("/login", response_class=HTMLResponse)
|
||||
async def login_post(request: Request, username: str = Form(...), password: str = Form(...)):
|
||||
@@ -54,9 +54,8 @@ async def login_post(request: Request, username: str = Form(...), password: str
|
||||
|
||||
if is_login_blocked(username, ip):
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="login.html",
|
||||
context={"error": "تم حجب محاولات تسجيل الدخول مؤقتاً. حاول بعد 15 دقيقة."},
|
||||
"login.html",
|
||||
{"request": request, "error": "تم حجب محاولات تسجيل الدخول مؤقتاً. حاول بعد 15 دقيقة."},
|
||||
status_code=429,
|
||||
)
|
||||
|
||||
@@ -64,9 +63,8 @@ async def login_post(request: Request, username: str = Form(...), password: str
|
||||
if not user or not verify_password(user["password_hash"], password):
|
||||
record_login_attempt(username, ip, success=False)
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="login.html",
|
||||
context={"error": "Invalid username or password"},
|
||||
"login.html",
|
||||
{"request": request, "error": "Invalid username or password"},
|
||||
status_code=401,
|
||||
)
|
||||
|
||||
@@ -109,9 +107,8 @@ async def logout(request: Request):
|
||||
async def users_list(request: Request, current=Depends(require_admin)):
|
||||
users = get_all_users()
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="users.html",
|
||||
context={"users": users, "current_user": current},
|
||||
"users.html",
|
||||
{"request": request, "users": users, "current_user": current},
|
||||
)
|
||||
|
||||
@router.post("/users/create")
|
||||
@@ -128,16 +125,14 @@ async def add_user(
|
||||
if len(password) < 8:
|
||||
users = get_all_users()
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="users.html",
|
||||
context={"users": users, "current_user": current, "error": "Password must be at least 8 characters."},
|
||||
"users.html",
|
||||
{"request": request, "users": users, "current_user": current, "error": "Password must be at least 8 characters."},
|
||||
)
|
||||
if get_user_by_username(username):
|
||||
users = get_all_users()
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="users.html",
|
||||
context={"users": users, "current_user": current, "error": f"User '{username}' already exists."},
|
||||
"users.html",
|
||||
{"request": request, "users": users, "current_user": current, "error": f"User '{username}' already exists."},
|
||||
)
|
||||
create_user(username, password, role=role)
|
||||
return RedirectResponse(url="/auth/users", status_code=status.HTTP_303_SEE_OTHER)
|
||||
@@ -147,9 +142,8 @@ async def remove_user(request: Request, user_id: int, current=Depends(require_ad
|
||||
if user_id == current["user_id"]:
|
||||
users = get_all_users()
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="users.html",
|
||||
context={"users": users, "current_user": current, "error": "لا يمكنك حذف حسابك الخاص."},
|
||||
"users.html",
|
||||
{"request": request, "users": users, "current_user": current, "error": "لا يمكنك حذف حسابك الخاص."},
|
||||
status_code=400,
|
||||
)
|
||||
# Prevent removing the last admin
|
||||
@@ -157,9 +151,8 @@ async def remove_user(request: Request, user_id: int, current=Depends(require_ad
|
||||
if target and target["role"] == "admin" and count_admins() <= 1:
|
||||
users = get_all_users()
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="users.html",
|
||||
context={"users": users, "current_user": current, "error": "لا يمكن حذف آخر مسؤول في النظام."},
|
||||
"users.html",
|
||||
{"request": request, "users": users, "current_user": current, "error": "لا يمكن حذف آخر مسؤول في النظام."},
|
||||
status_code=400,
|
||||
)
|
||||
delete_user(user_id)
|
||||
@@ -180,9 +173,8 @@ async def change_role(
|
||||
if target and target["role"] == "admin" and role != "admin" and count_admins() <= 1:
|
||||
users = get_all_users()
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="users.html",
|
||||
context={"users": users, "current_user": current, "error": "لا يمكن تخفيض رتبة آخر مسؤول."},
|
||||
"users.html",
|
||||
{"request": request, "users": users, "current_user": current, "error": "لا يمكن تخفيض رتبة آخر مسؤول."},
|
||||
status_code=400,
|
||||
)
|
||||
set_user_role(user_id, role)
|
||||
@@ -196,9 +188,9 @@ async def backup_db(request: Request, current=Depends(require_admin)):
|
||||
except Exception as e:
|
||||
msg = f"Backup failed: {str(e)}"
|
||||
return templates.TemplateResponse(
|
||||
request=request,
|
||||
name="users.html",
|
||||
context={
|
||||
"users.html",
|
||||
{
|
||||
"request": request,
|
||||
"users": get_all_users(),
|
||||
"current_user": current,
|
||||
"backup_msg": msg,
|
||||
|
||||
+3
-1
@@ -12,6 +12,7 @@ from services.extractor import (
|
||||
extract_document,
|
||||
get_available_providers,
|
||||
get_default_provider,
|
||||
provider_supports_ai_verification,
|
||||
verify_page_correlation,
|
||||
)
|
||||
from services.search_service import normalize_arabic, _normalize_scope
|
||||
@@ -336,7 +337,8 @@ async def review_document(request: Request, doc_id: int, wait: int = 0):
|
||||
"providers": get_available_providers(),
|
||||
"current_provider": doc.get("provider") or get_default_provider(),
|
||||
"ai_verification_available": any(
|
||||
provider["id"] in {"claude", "gemini"} for provider in get_available_providers()
|
||||
provider_supports_ai_verification(provider["id"])
|
||||
for provider in get_available_providers()
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
+94
-83
@@ -9,7 +9,7 @@ from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile
|
||||
from fastapi.responses import JSONResponse, RedirectResponse
|
||||
from fastapi.templating import Jinja2Templates
|
||||
|
||||
from config import UPLOAD_DIR
|
||||
from config import UPLOAD_DIR, MAX_CONCURRENT_EXTRACTIONS
|
||||
from database.connection import get_db
|
||||
from services.extractor import extract_document, get_available_providers, get_default_provider
|
||||
from services.pdf_handler import pdf_to_images
|
||||
@@ -17,6 +17,16 @@ from services.pdf_handler import pdf_to_images
|
||||
router = APIRouter()
|
||||
templates = Jinja2Templates(directory="templates")
|
||||
|
||||
# Semaphore: limits simultaneous AI API calls to avoid quota exhaustion on bulk uploads
|
||||
_extraction_semaphore: asyncio.Semaphore | None = None
|
||||
|
||||
|
||||
def _get_semaphore() -> asyncio.Semaphore:
|
||||
global _extraction_semaphore
|
||||
if _extraction_semaphore is None:
|
||||
_extraction_semaphore = asyncio.Semaphore(MAX_CONCURRENT_EXTRACTIONS)
|
||||
return _extraction_semaphore
|
||||
|
||||
ALLOWED_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"}
|
||||
ALLOWED_PDF_EXTS = {".pdf"}
|
||||
ALLOWED_EXTENSIONS = ALLOWED_IMAGE_EXTS | ALLOWED_PDF_EXTS
|
||||
@@ -91,95 +101,96 @@ def _save_image(file_bytes: bytes, original_name: str) -> str:
|
||||
|
||||
async def _extract_and_save(doc_id: int, image_path: str, provider: str = ""):
|
||||
"""Background task: call extractor, parse result, update DB."""
|
||||
try:
|
||||
data = await extract_document(image_path, provider=provider)
|
||||
raw_json = json.dumps(data, ensure_ascii=False)
|
||||
async with _get_semaphore():
|
||||
try:
|
||||
data = await extract_document(image_path, provider=provider)
|
||||
raw_json = json.dumps(data, ensure_ascii=False)
|
||||
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"""UPDATE documents SET status='extracted',
|
||||
raw_extraction_json=?,
|
||||
request_number=?,
|
||||
request_date=?,
|
||||
applicant_name_raw=?,
|
||||
request_purpose=?,
|
||||
data_valid_until=?,
|
||||
registry_office=?,
|
||||
owns_properties=?,
|
||||
declared_property_count=?,
|
||||
page_info=?,
|
||||
search_scope=?,
|
||||
updated_at=CURRENT_TIMESTAMP
|
||||
WHERE id=?""",
|
||||
(
|
||||
raw_json,
|
||||
data.get("request_number"),
|
||||
data.get("request_date"),
|
||||
data.get("applicant_name_raw"),
|
||||
data.get("request_purpose"),
|
||||
data.get("data_valid_until"),
|
||||
data.get("registry_office"),
|
||||
data.get("owns_properties"),
|
||||
data.get("declared_property_count"),
|
||||
data.get("page_info"),
|
||||
data.get("search_scope"),
|
||||
doc_id,
|
||||
),
|
||||
)
|
||||
|
||||
for i, prop in enumerate(data.get("properties", [])):
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"""INSERT INTO properties
|
||||
(document_id, row_order, party_name, property_number,
|
||||
section, block, real_estate_district, qaza, num_shares, ownership_type)
|
||||
VALUES (?,?,?,?,?,?,?,?,?,?)""",
|
||||
"""UPDATE documents SET status='extracted',
|
||||
raw_extraction_json=?,
|
||||
request_number=?,
|
||||
request_date=?,
|
||||
applicant_name_raw=?,
|
||||
request_purpose=?,
|
||||
data_valid_until=?,
|
||||
registry_office=?,
|
||||
owns_properties=?,
|
||||
declared_property_count=?,
|
||||
page_info=?,
|
||||
search_scope=?,
|
||||
updated_at=CURRENT_TIMESTAMP
|
||||
WHERE id=?""",
|
||||
(
|
||||
doc_id, i,
|
||||
prop.get("party_name"),
|
||||
prop.get("property_number"),
|
||||
prop.get("section"),
|
||||
prop.get("block"),
|
||||
prop.get("real_estate_district"),
|
||||
prop.get("qaza"),
|
||||
prop.get("num_shares"),
|
||||
prop.get("ownership_type"),
|
||||
raw_json,
|
||||
data.get("request_number"),
|
||||
data.get("request_date"),
|
||||
data.get("applicant_name_raw"),
|
||||
data.get("request_purpose"),
|
||||
data.get("data_valid_until"),
|
||||
data.get("registry_office"),
|
||||
data.get("owns_properties"),
|
||||
data.get("declared_property_count"),
|
||||
data.get("page_info"),
|
||||
data.get("search_scope"),
|
||||
doc_id,
|
||||
),
|
||||
)
|
||||
|
||||
# Flag logical duplicates: same request_number + search_scope + page_number
|
||||
req_num = (data.get("request_number") or "").strip()
|
||||
scope = (data.get("search_scope") or "").strip()
|
||||
page_info = (data.get("page_info") or "").strip()
|
||||
if req_num:
|
||||
existing = conn.execute(
|
||||
"""SELECT id FROM documents
|
||||
WHERE id != ? AND request_number=?
|
||||
AND COALESCE(search_scope,'')=?
|
||||
AND COALESCE(page_info,'')=?
|
||||
AND status IN ('extracted','confirmed')
|
||||
AND COALESCE(duplicate_dismissed, 0) = 0
|
||||
ORDER BY id LIMIT 1""",
|
||||
(doc_id, req_num, scope, page_info),
|
||||
).fetchone()
|
||||
if existing:
|
||||
# Only auto-flag if THIS document hasn't itself been dismissed.
|
||||
self_row = conn.execute(
|
||||
"SELECT COALESCE(duplicate_dismissed, 0) AS d FROM documents WHERE id=?",
|
||||
(doc_id,),
|
||||
).fetchone()
|
||||
if not (self_row and self_row["d"]):
|
||||
conn.execute(
|
||||
"UPDATE documents SET duplicate_of=? WHERE id=?",
|
||||
(existing["id"], doc_id),
|
||||
)
|
||||
for i, prop in enumerate(data.get("properties", [])):
|
||||
conn.execute(
|
||||
"""INSERT INTO properties
|
||||
(document_id, row_order, party_name, property_number,
|
||||
section, block, real_estate_district, qaza, num_shares, ownership_type)
|
||||
VALUES (?,?,?,?,?,?,?,?,?,?)""",
|
||||
(
|
||||
doc_id, i,
|
||||
prop.get("party_name"),
|
||||
prop.get("property_number"),
|
||||
prop.get("section"),
|
||||
prop.get("block"),
|
||||
prop.get("real_estate_district"),
|
||||
prop.get("qaza"),
|
||||
prop.get("num_shares"),
|
||||
prop.get("ownership_type"),
|
||||
),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"""UPDATE documents SET status='error', extraction_error=?,
|
||||
updated_at=CURRENT_TIMESTAMP WHERE id=?""",
|
||||
(str(e), doc_id),
|
||||
)
|
||||
# Flag logical duplicates: same request_number + search_scope + page_number
|
||||
req_num = (data.get("request_number") or "").strip()
|
||||
scope = (data.get("search_scope") or "").strip()
|
||||
page_info = (data.get("page_info") or "").strip()
|
||||
if req_num:
|
||||
existing = conn.execute(
|
||||
"""SELECT id FROM documents
|
||||
WHERE id != ? AND request_number=?
|
||||
AND COALESCE(search_scope,'')=?
|
||||
AND COALESCE(page_info,'')=?
|
||||
AND status IN ('extracted','confirmed')
|
||||
AND COALESCE(duplicate_dismissed, 0) = 0
|
||||
ORDER BY id LIMIT 1""",
|
||||
(doc_id, req_num, scope, page_info),
|
||||
).fetchone()
|
||||
if existing:
|
||||
# Only auto-flag if THIS document hasn't itself been dismissed.
|
||||
self_row = conn.execute(
|
||||
"SELECT COALESCE(duplicate_dismissed, 0) AS d FROM documents WHERE id=?",
|
||||
(doc_id,),
|
||||
).fetchone()
|
||||
if not (self_row and self_row["d"]):
|
||||
conn.execute(
|
||||
"UPDATE documents SET duplicate_of=? WHERE id=?",
|
||||
(existing["id"], doc_id),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
with get_db() as conn:
|
||||
conn.execute(
|
||||
"""UPDATE documents SET status='error', extraction_error=?,
|
||||
updated_at=CURRENT_TIMESTAMP WHERE id=?""",
|
||||
(str(e), doc_id),
|
||||
)
|
||||
|
||||
|
||||
@router.get("/")
|
||||
|
||||
+50
-14
@@ -1,4 +1,5 @@
|
||||
import base64
|
||||
import importlib.util
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
@@ -61,7 +62,7 @@ Rules:
|
||||
|
||||
def _get_ai_verification_provider(preferred_provider: str = "") -> str:
|
||||
"""Return a provider capable of vision reasoning for verification."""
|
||||
providers = [p["id"] for p in get_available_providers() if p["id"] in {"claude", "gemini"}]
|
||||
providers = [p["id"] for p in get_available_providers() if provider_supports_ai_verification(p["id"])]
|
||||
if preferred_provider in providers:
|
||||
return preferred_provider
|
||||
if DEFAULT_PROVIDER in providers:
|
||||
@@ -71,6 +72,10 @@ def _get_ai_verification_provider(preferred_provider: str = "") -> str:
|
||||
raise ValueError("No AI verification provider configured")
|
||||
|
||||
|
||||
def provider_supports_ai_verification(provider_id: str) -> bool:
|
||||
return provider_id == "claude" or provider_id == "gemini" or provider_id.startswith("gemini-")
|
||||
|
||||
|
||||
def _build_correlation_user_prompt(current_context: dict, candidate_context: dict) -> str:
|
||||
return (
|
||||
"Determine whether PAGE_A and PAGE_B belong to the same multi-page request/document for the same person. "
|
||||
@@ -249,15 +254,29 @@ def _parse_json_lenient(text: str) -> dict:
|
||||
raise ValueError(f"Could not parse model response as JSON: {text[:200]}")
|
||||
|
||||
|
||||
# Individual Gemini models exposed to the UI
|
||||
_GEMINI_MODEL_OPTIONS = [
|
||||
("gemini-2.5-pro", "Gemini 2.5 Pro (أفضل دقة)"),
|
||||
("gemini-2.5-flash", "Gemini 2.5 Flash (سريع)"),
|
||||
("gemini-2.0-flash", "Gemini 2.0 Flash"),
|
||||
("gemini-2.0-flash-lite", "Gemini 2.0 Flash Lite (احتياطي)"),
|
||||
]
|
||||
|
||||
|
||||
def _easyocr_is_available() -> bool:
|
||||
return importlib.util.find_spec("easyocr") is not None
|
||||
|
||||
|
||||
def get_available_providers() -> list[dict]:
|
||||
"""Return list of all available providers (EasyOCR is always available)."""
|
||||
providers = [
|
||||
{"id": "easyocr", "name": "EasyOCR (مجاني)", "model": "local"},
|
||||
]
|
||||
"""Return only providers that are actually usable in the current runtime."""
|
||||
providers = []
|
||||
if _easyocr_is_available():
|
||||
providers.append({"id": "easyocr", "name": "EasyOCR (مجاني)", "model": "local"})
|
||||
if ANTHROPIC_API_KEY:
|
||||
providers.append({"id": "claude", "name": "Claude (Anthropic)", "model": CLAUDE_MODEL})
|
||||
if GEMINI_API_KEY:
|
||||
providers.append({"id": "gemini", "name": "Gemini (Google)", "model": GEMINI_MODEL})
|
||||
for model_id, label in _GEMINI_MODEL_OPTIONS:
|
||||
providers.append({"id": model_id, "name": label, "model": model_id})
|
||||
return providers
|
||||
|
||||
|
||||
@@ -265,8 +284,12 @@ def get_default_provider() -> str:
|
||||
"""Return the default provider, falling back to whichever is available."""
|
||||
providers = get_available_providers()
|
||||
if not providers:
|
||||
return "easyocr"
|
||||
return GEMINI_MODEL if DEFAULT_PROVIDER == "gemini" else DEFAULT_PROVIDER
|
||||
ids = [p["id"] for p in providers]
|
||||
# Map legacy "gemini" default to the configured GEMINI_MODEL
|
||||
resolved = GEMINI_MODEL if DEFAULT_PROVIDER == "gemini" else DEFAULT_PROVIDER
|
||||
if resolved in ids:
|
||||
return resolved
|
||||
if DEFAULT_PROVIDER in ids:
|
||||
return DEFAULT_PROVIDER
|
||||
return ids[0]
|
||||
@@ -329,7 +352,7 @@ async def _extract_with_claude(image_path: str) -> dict:
|
||||
|
||||
# ─── Gemini extraction ───────────────────────────────────────────
|
||||
|
||||
async def _extract_with_gemini(image_path: str) -> dict:
|
||||
async def _extract_with_gemini(image_path: str, model: str = "") -> dict:
|
||||
import asyncio
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
@@ -345,8 +368,17 @@ async def _extract_with_gemini(image_path: str) -> dict:
|
||||
".png": "image/png", ".webp": "image/webp"}
|
||||
mime_type = mime_map.get(suffix, "image/jpeg")
|
||||
|
||||
# Try primary model first; only fall back on 503 (overloaded), NOT on 429 (quota)
|
||||
models_to_try = [GEMINI_MODEL, "gemini-2.0-flash-lite"]
|
||||
# Full fallback chain: start from requested model, cascade through cheaper/available ones
|
||||
_ALL_GEMINI_FALLBACK = [
|
||||
"gemini-2.5-pro",
|
||||
"gemini-2.5-flash",
|
||||
"gemini-2.0-flash",
|
||||
"gemini-2.0-flash-lite",
|
||||
]
|
||||
# If a specific model was requested, start from it; otherwise start from the configured default
|
||||
start_model = model if model else GEMINI_MODEL
|
||||
# Build ordered list: requested model first, then remaining fallbacks in order
|
||||
models_to_try = [start_model] + [m for m in _ALL_GEMINI_FALLBACK if m != start_model]
|
||||
|
||||
def _call(model_name: str):
|
||||
response = client.models.generate_content(
|
||||
@@ -731,6 +763,8 @@ async def extract_document(image_path: str, provider: str = "") -> dict:
|
||||
"""
|
||||
Extract structured data from a document image using the specified provider.
|
||||
Falls back to the default provider if none specified.
|
||||
provider can be: "claude", "easyocr", "gemini" (uses GEMINI_MODEL),
|
||||
or a specific Gemini model ID like "gemini-2.5-flash".
|
||||
"""
|
||||
if not provider:
|
||||
provider = get_default_provider()
|
||||
@@ -739,11 +773,13 @@ async def extract_document(image_path: str, provider: str = "") -> dict:
|
||||
if not ANTHROPIC_API_KEY:
|
||||
raise ValueError("ANTHROPIC_API_KEY not set")
|
||||
return await _extract_with_claude(image_path)
|
||||
elif provider == "gemini":
|
||||
if not GEMINI_API_KEY:
|
||||
raise ValueError("GEMINI_API_KEY not set")
|
||||
return await _extract_with_gemini(image_path)
|
||||
elif provider == "easyocr":
|
||||
return await _extract_with_easyocr(image_path)
|
||||
elif provider == "gemini" or provider.startswith("gemini-"):
|
||||
if not GEMINI_API_KEY:
|
||||
raise ValueError("GEMINI_API_KEY not set")
|
||||
# Pass the specific model if the provider ID encodes one
|
||||
model_override = provider if provider.startswith("gemini-") else ""
|
||||
return await _extract_with_gemini(image_path, model=model_override)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
+15
-4
@@ -6,25 +6,35 @@ from pathlib import Path
|
||||
|
||||
from config import UPLOAD_DIR
|
||||
|
||||
# Maximum pages we will extract from a single PDF upload.
|
||||
# Beyond this limit pages are silently dropped to prevent runaway queue growth.
|
||||
MAX_PDF_PAGES = 200
|
||||
|
||||
|
||||
def pdf_to_images(pdf_bytes: bytes, original_name: str) -> list[dict]:
|
||||
"""
|
||||
Convert a PDF to individual page images.
|
||||
Returns list of dicts: [{"image_path": "relative/path.png", "page_number": 1}, ...]
|
||||
|
||||
Each page is rendered at 250 DPI — sufficient for Arabic OCR while keeping
|
||||
file sizes manageable. For scanned/image-only PDFs PyMuPDF simply renders
|
||||
the embedded raster; no additional image extraction step is needed.
|
||||
|
||||
Returns list of dicts: [{"image_path": "relative/path.png", "page_number": 1, ...}, ...]
|
||||
"""
|
||||
import fitz # PyMuPDF
|
||||
|
||||
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
|
||||
total_pages = len(doc)
|
||||
group_id = uuid.uuid4().hex
|
||||
today = date.today().isoformat()
|
||||
dest_dir = Path(UPLOAD_DIR) / today
|
||||
dest_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
pages = []
|
||||
for page_num in range(len(doc)):
|
||||
for page_num in range(min(total_pages, MAX_PDF_PAGES)):
|
||||
page = doc[page_num]
|
||||
# Render at 200 DPI for good OCR quality
|
||||
pix = page.get_pixmap(dpi=200)
|
||||
# 250 DPI gives sharper Arabic text while staying under ~2 MB per PNG
|
||||
pix = page.get_pixmap(dpi=250)
|
||||
|
||||
filename = f"{group_id}_p{page_num + 1}.png"
|
||||
dest = dest_dir / filename
|
||||
@@ -34,6 +44,7 @@ def pdf_to_images(pdf_bytes: bytes, original_name: str) -> list[dict]:
|
||||
"image_path": str(Path(today) / filename),
|
||||
"page_number": page_num + 1,
|
||||
"pdf_group_id": group_id,
|
||||
"total_pages": total_pages,
|
||||
})
|
||||
|
||||
doc.close()
|
||||
|
||||
Reference in New Issue
Block a user