Spaces:
Build error
Build error
File size: 4,870 Bytes
cd964f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import sys
import os
import uuid
import tempfile
import pandas as pd
import torch
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from typing import List, Optional
from src.extractor import PassportExtractor
from src.formats import format_fly_dubai, format_iraqi, format_fly_baghdad
# Path fix for local imports
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
app = FastAPI()
# Initialize extractor
use_gpu = torch.cuda.is_available()
extractor = PassportExtractor(use_gpu=use_gpu)
# Temporary storage for download files
DOWNLOAD_DIR = os.path.join(tempfile.gettempdir(), "passport_downloads")
os.makedirs(DOWNLOAD_DIR, exist_ok=True)
@app.post("/api/process")
async def process_passports(
files: List[UploadFile] = File(...),
airline: str = Form(...),
include_all: bool = Form(False)
):
try:
all_results = []
problematic_files = []
# Save uploaded files temporarily to process
for file in files:
file_name = file.filename
content = await file.read()
with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file_name)[1]) as tmp:
tmp.write(content)
tmp_path = tmp.name
try:
is_pdf = file_name.lower().endswith('.pdf')
if is_pdf:
results = extractor.process_pdf(tmp_path, airline=airline.lower())
if not results:
results = [{"mrz_found": False}] # Placeholder for failure
else:
result = extractor.get_data(tmp_path, airline=airline.lower())
results = [result] if result else []
# Attach source file name
mrz_found_in_file = False
for res in results:
res['source_file'] = file_name
if res.get('mrz_found', False):
mrz_found_in_file = True
all_results.extend(results)
if not results or not mrz_found_in_file:
problematic_files.append({'file_name': file_name, 'reason': 'MRZ not found'})
except Exception as e:
problematic_files.append({'file_name': file_name, 'reason': str(e)})
finally:
if os.path.exists(tmp_path):
os.remove(tmp_path)
# Filter and format
good_results = [r for r in all_results if not any(p['file_name'] == r.get('source_file') for p in problematic_files)]
if airline == "Fly Dubai":
df = format_fly_dubai(good_results)
elif airline == "Iraqi":
df = format_iraqi(good_results)
elif airline == "Fly Baghdad":
df = format_fly_baghdad(good_results)
else:
df = pd.DataFrame(good_results)
# Add placeholders if requested
if include_all and problematic_files:
# Re-implementing simplified placeholder logic for API
placeholders = []
for p in problematic_files:
placeholders.append({"Last Name": "•••", "First Name": "•••", "source_file": p['file_name']})
df = pd.concat([df, pd.DataFrame(placeholders)], ignore_index=True)
# Generate unique IDs for downloads
session_id = str(uuid.uuid4())
csv_name = f"passports_{session_id}.csv"
xlsx_name = f"passports_{session_id}.xlsx"
df.to_csv(os.path.join(DOWNLOAD_DIR, csv_name), index=False)
df.to_excel(os.path.join(DOWNLOAD_DIR, xlsx_name), index=False)
return {
"data": df.fillna("").to_dict(orient="records"),
"csv_id": csv_name,
"xlsx_id": xlsx_name,
"summary": {
"total": len(files),
"success": len(good_results),
"failed": len(problematic_files)
}
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/download/{file_name}")
async def download_file(file_name: str):
file_path = os.path.join(DOWNLOAD_DIR, file_name)
if os.path.exists(file_path):
return FileResponse(file_path, filename=file_name)
raise HTTPException(status_code=404, detail="File not found")
# Serve static files (HTML/JS)
app.mount("/", StaticFiles(directory="static", html=True), name="static")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
|