app.py

import os
import sys
import tempfile
from pathlib import Path
import pandas as pd

# Keep Gradio on client-side rendering. In hosted Gradio 6 runs, SSR starts a
# Node proxy that can emit noisy asyncio cleanup tracebacks at startup/shutdown.
os.environ.setdefault("GRADIO_SSR_MODE", "false")

import gradio as gr
import time

# Ensure project root is in python path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

try:
    import spaces
except ImportError:
    spaces = None

from extractor import (
    PassportExtractor,
    format_fly_baghdad,
    format_fly_dubai,
    format_iraqi,
)


def gpu_task(duration=120):
    if spaces is None:
        return lambda fn: fn
    return spaces.GPU(duration=duration)


# Initialize Extractor
print("Initializing Passport Extractor...")
EXTRACTOR = PassportExtractor(use_gpu=False)
print("Extractor initialized.")


MAX_FILE_SIZE_MB = 100
ALLOWED_FILE_TYPES = [
    ".png",
    ".jpg",
    ".jpeg",
    ".pdf",
    ".avif",
    ".webp",
    ".bmp",
    ".tiff",
]


def _format_data(good_results, airline):
    if airline == "Fly Dubai":
        return format_fly_dubai(good_results)
    if airline == "Iraqi":
        return format_iraqi(good_results)
    if airline == "Fly Baghdad":
        return format_fly_baghdad(good_results)

    # Default view
    df = pd.DataFrame(good_results)
    if not df.empty:
        from extractor import format_date
        if "date_of_birth" in df.columns:
            df["date_of_birth"] = df["date_of_birth"].apply(lambda x: format_date(x, "%d/%m/%Y", is_dob=True))
        if "expiration_date" in df.columns:
            df["expiration_date"] = df["expiration_date"].apply(lambda x: format_date(x, "%d/%m/%Y", is_dob=False))
    return df


@gpu_task(duration=300)
def process_files(files, airline, export_all_files, progress=gr.Progress()):
    if not files:
        return None, "No files uploaded.", None, None
    start_time = time.time()
    all_results, problematic_files = [], []
    total_files = len(files)
    from concurrent.futures import ThreadPoolExecutor, as_completed

    def _uploaded_path(file_obj):
        if isinstance(file_obj, (str, os.PathLike)):
            return os.fspath(file_obj)
        if hasattr(file_obj, "name"):
            return file_obj.name
        if isinstance(file_obj, dict):
            return file_obj.get("path") or file_obj.get("name")
        return None

    def process_single_file(file_obj):
        file_path = _uploaded_path(file_obj)
        if not file_path:
            return [], {
                "file_name": "unknown",
                "reason": f"Unsupported upload object: {type(file_obj).__name__}",
            }

        file_name = Path(file_path).name
        ext = os.path.splitext(file_name)[1].lower()
        file_results, file_problem = [], None
        if os.path.getsize(file_path) > MAX_FILE_SIZE_MB * 1024 * 1024:
            return [], {
                "file_name": file_name,
                "reason": f"File too large. Max {MAX_FILE_SIZE_MB} MB allowed.",
            }
        try:
            if ext == ".pdf":

                def pdf_progress(p):
                    progress(
                        completed_files / total_files + (p * (1 / total_files)),
                        desc=f"Processing {file_name}...",
                    )

                file_results = (
                    EXTRACTOR.process_pdf(
                        file_path,
                        progress_callback=pdf_progress,
                        airline=(airline or "Default").lower(),
                    )
                    or []
                )
            else:
                result = EXTRACTOR.get_data(
                    file_path, airline=(airline or "Default").lower()
                )
                if result:
                    file_results = [result]
        except Exception as e:
            file_problem = {"file_name": file_name, "reason": f"Error: {e}"}
        mrz_found = False
        for res in file_results:
            res["source_file"] = file_name
            if res.get("mrz_found"):
                mrz_found = True
        if not file_problem:
            if not file_results:
                file_problem = {
                    "file_name": file_name,
                    "reason": "No passport data detected",
                }
            elif not mrz_found:
                file_problem = {"file_name": file_name, "reason": "MRZ data not found"}
        return file_results, file_problem

    completed_files = 0
    with ThreadPoolExecutor(max_workers=2) as executor:
        future_to_file = {executor.submit(process_single_file, f): f for f in files}
        for future in as_completed(future_to_file):
            completed_files += 1
            try:
                file_results, file_problem = future.result()
            except Exception as e:
                original_path = _uploaded_path(future_to_file[future]) or "unknown"
                file_results = []
                file_problem = {
                    "file_name": Path(original_path).name,
                    "reason": f"Error: {e}",
                }
            all_results.extend(file_results)
            if file_problem:
                problematic_files.append(file_problem)

    good_results = [r for r in all_results if r.get("mrz_found")]
    end_time = time.time()
    duration = round(end_time - start_time, 2)

    if not good_results and problematic_files:
        status_msg = f"No data could be extracted. (Time: {duration}s)"
        status_msg += "\n\nProblems:\n" + "\n".join(
            [f"- {p['file_name']}: {p['reason']}" for p in problematic_files]
        )
        return None, status_msg, None, None

    if not good_results:
        return (
            None,
            f"⚠️ No data could be extracted. (Time: {duration}s)",
            None,
            None,
        )

    df = _format_data(good_results, airline)
    status_msg = f"Extracted {len(good_results)} passport(s) from {total_files} file(s) in {duration} seconds."
    if problematic_files:
        status_msg += "\n\nProblems:\n" + "\n".join(
            [f"- {p['file_name']}: {p['reason']}" for p in problematic_files]
        )

    csv_path = tempfile.mktemp(suffix=".csv")
    df.to_csv(csv_path, index=False)
    excel_path = tempfile.mktemp(suffix=".xlsx")
    df.to_excel(excel_path, index=False)
    return df, status_msg, csv_path, excel_path


def main():
    with gr.Blocks(title="Passport OCR Extractor") as demo:
        gr.Markdown("#Passport OCR Extractor")
        with gr.Row():
            with gr.Column(scale=1):
                airline = gr.Dropdown(
                    choices=["Default", "Fly Dubai", "Iraqi", "Fly Baghdad"],
                    value="Default",
                    label="Airline Format",
                )
                file_input = gr.File(
                    file_count="multiple",
                    label="Upload Passports",
                    file_types=ALLOWED_FILE_TYPES,
                )
                extract_btn = gr.Button("Extract Data", variant="primary")
            with gr.Column(scale=2):
                status_output = gr.Markdown()
                table_output = gr.Dataframe(label="Extracted Data")
                with gr.Row():
                    csv_download = gr.File(label="Download CSV")
                    excel_download = gr.File(label="Download Excel")
        extract_btn.click(
            fn=process_files,
            inputs=[file_input, airline, gr.State(False)],
            outputs=[table_output, status_output, csv_download, excel_download],
        )
    port = int(os.environ.get("PORT", "7860"))
    demo.launch(server_name="0.0.0.0", server_port=port, show_error=True, ssr_mode=False)


if __name__ == "__main__":
    main()