Spaces:
Sleeping
Sleeping
Gemini CLI
Final stability and formatting fixes: corrected templates, strict OCR, and pagination
90ad37b | import os | |
| import sys | |
| import tempfile | |
| from io import BytesIO | |
| import streamlit as st | |
| import pandas as pd | |
| from datetime import datetime | |
| # Ensure project root is in python path | |
| sys.path.append(os.path.dirname(os.path.abspath(__file__))) | |
| from src.extractor import PassportExtractor | |
| from src.formats import format_fly_baghdad, format_fly_dubai, format_iraqi | |
| # Page Config | |
| st.set_page_config(page_title="Passport OCR Extractor", page_icon="🛂", layout="wide") | |
| # Custom CSS for UI polish | |
| st.markdown(""" | |
| <style> | |
| /* Change cursor to pointer for selectbox */ | |
| div[data-baseweb="select"] { | |
| cursor: pointer !important; | |
| } | |
| /* Hide the upload file list while processing if needed, | |
| but Streamlit handles this naturally with spinners */ | |
| </style> | |
| """, unsafe_allow_html=True) | |
| # Initialize Extractor | |
| def get_extractor(): | |
| return PassportExtractor(use_gpu=True) | |
| EXTRACTOR = get_extractor() | |
| MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024 | |
| ALLOWED_FILE_TYPES = ["png", "jpg", "jpeg", "pdf", "avif", "webp", "bmp", "tiff"] | |
| def _safe_reason(problem): | |
| return ( | |
| problem.get("reason") | |
| or problem.get("issue") | |
| or problem.get("message") | |
| or problem.get("error") | |
| or "<unspecified>" | |
| ) | |
| 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) | |
| return pd.DataFrame(good_results) | |
| def _placeholder_rows(airline, problematic_files, good_results): | |
| rows = [] | |
| for problem in problematic_files: | |
| if airline == "Fly Dubai": | |
| rows.append({ | |
| "TYPE": "Adult", "TITLE": "MR", "FIRST NAME": "•••", | |
| "LAST NAME": "•••", "DOB (DDMMMYY)": "•••", "GENDER": "Male" | |
| }) | |
| elif airline == "Iraqi": | |
| rows.append({ | |
| "Last Name": "•••", "First Name and Middle Name": "•••", "Title": "MR", | |
| "PTC": "ADT", "Gender": "M", "Date of Birth": "•••", | |
| "Passport Last Name": "•••", "Passport First Name": "•••", | |
| "Passport Middle Name": "", "Passport Number": "•••", | |
| "Passport Nationality": "•••", "Passport Issue Country": "•••", | |
| "Passport Expiry Date": "•••", "Visa Number": "", "Visa Type": "", | |
| "Visa Issue Date": "", "Place of Birth": "", "Visa Place of Issue": "", | |
| "Visa Country of Application": "", "Address Type": "", "Address Country": "", | |
| "Address Details": "", "Address City": "", "Address State": "", "Address Zip Code": "" | |
| }) | |
| elif airline == "Fly Baghdad": | |
| rows.append({ | |
| "Sequence": len(good_results) + len(rows) + 1, "Traveling With": "", | |
| "Pax Type": "ADT", "Title": "MR", "First Name": "•••", "Last Name": "•••", | |
| "Gender": "MALE", "DOB (dd/mm/yyyy)": "•••", "Nationality": "•••", | |
| "Passport Number": "•••", "Passport Expiry (dd/mm/yyyy)": "•••", | |
| "Passport Issued Country": "•••" | |
| }) | |
| else: | |
| rows.append({ | |
| "source_file": problem["file_name"], "surname": "•••", "given_names": "•••", | |
| "passport_number": "•••", "nationality": "•••", "date_of_birth": "•••", | |
| "sex": "•••", "expiration_date": "•••", "personal_number": "•••", "mrz_found": False | |
| }) | |
| return rows | |
| def main(): | |
| st.title("🛂 Passport OCR Extractor") | |
| st.markdown("Extract data from passport MRZ codes into structured tables.") | |
| with st.sidebar: | |
| st.header("Settings") | |
| st.info("💡 Optimized for performance - GPU enabled") | |
| airline = st.selectbox( | |
| "Choose Airline Format", | |
| options=["Default", "Fly Dubai", "Iraqi", "Fly Baghdad"], | |
| index=0 | |
| ) | |
| export_all_files = st.checkbox( | |
| "Include all files in export", | |
| value=False, | |
| help="When checked, includes files with missing/invalid data (shows ••• for missing fields)" | |
| ) | |
| uploaded_files = st.file_uploader( | |
| "Upload Passport Files", | |
| type=ALLOWED_FILE_TYPES, | |
| accept_multiple_files=True | |
| ) | |
| if st.button("Extract Data") and uploaded_files: | |
| all_results = [] | |
| problematic_files = [] | |
| # UI placeholder for results - cleared on each run | |
| status_container = st.empty() | |
| table_container = st.empty() | |
| download_container = st.empty() | |
| failed_container = st.expander("📋 Failed passport details", expanded=False) | |
| with st.status("Processing files...", expanded=True) as status: | |
| total_files = len(uploaded_files) | |
| for i, uploaded_file in enumerate(uploaded_files): | |
| file_name = uploaded_file.name | |
| ext = os.path.splitext(file_name)[1].lower() | |
| status.update(label=f"Processing {i+1}/{total_files}: {file_name}") | |
| if uploaded_file.size > MAX_FILE_SIZE_BYTES: | |
| problematic_files.append({ | |
| "file_name": file_name, | |
| "reason": f"File too large. Max 10 MB allowed." | |
| }) | |
| continue | |
| # Save temp file for processing | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_file: | |
| tmp_file.write(uploaded_file.getvalue()) | |
| tmp_path = tmp_file.name | |
| file_processed = False | |
| results = [] | |
| try: | |
| if ext == ".pdf": | |
| results = EXTRACTOR.process_pdf(tmp_path, airline=airline.lower()) or [] | |
| file_processed = True | |
| else: | |
| result = EXTRACTOR.get_data(tmp_path, airline=airline.lower()) | |
| if result: | |
| results = [result] | |
| file_processed = True | |
| except Exception as e: | |
| problematic_files.append({"file_name": file_name, "reason": f"Processing error: {e}"}) | |
| finally: | |
| if os.path.exists(tmp_path): | |
| os.remove(tmp_path) | |
| 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 file_processed: | |
| problematic_files.append({"file_name": file_name, "reason": "Unsupported format or corrupted file"}) | |
| elif not results: | |
| problematic_files.append({"file_name": file_name, "reason": "No passport data detected"}) | |
| elif not mrz_found_in_file: | |
| problematic_files.append({"file_name": file_name, "reason": "MRZ data not found"}) | |
| status.update(label="Extraction Complete!", state="complete", expanded=False) | |
| # Filtering logic | |
| good_results = [] | |
| for result in all_results: | |
| is_problematic = any(p["file_name"] == result.get("source_file") for p in problematic_files) | |
| if not is_problematic and result.get("mrz_found", False): | |
| good_results.append(result) | |
| successful_extractions = len(good_results) | |
| failed_count = total_files - len(set(r.get("source_file", "") for r in good_results)) | |
| if failed_count > 0: | |
| status_container.warning(f"⚠️ {successful_extractions} passport(s) were successfully extracted from {total_files} file(s).") | |
| else: | |
| status_container.success(f"✅ All {successful_extractions} passport(s) were successfully imported.") | |
| if problematic_files: | |
| with failed_container: | |
| for problem in problematic_files: | |
| st.write(f"- **{problem['file_name']}**: {_safe_reason(problem)}") | |
| if not good_results and not (export_all_files and problematic_files): | |
| st.error("⚠️ No data could be extracted. Please check the files.") | |
| return | |
| # Prepare final DataFrame | |
| df = _format_data(good_results, airline) | |
| if export_all_files and problematic_files: | |
| placeholders = _placeholder_rows(airline, problematic_files, good_results) | |
| if placeholders: | |
| placeholder_df = pd.DataFrame(placeholders) | |
| df = placeholder_df if df.empty else pd.concat([df, placeholder_df], ignore_index=True) | |
| # Display results with pagination | |
| table_container.subheader(f"📊 Extracted Data ({len(df)} records)") | |
| # Streamlit's st.dataframe handles pagination automatically in recent versions | |
| table_container.dataframe(df, use_container_width=True, hide_index=True) | |
| # Download buttons | |
| col1, col2 = download_container.columns(2) | |
| csv = df.to_csv(index=False).encode('utf-8') | |
| col1.download_button( | |
| label="Download data as CSV", | |
| data=csv, | |
| file_name=f"passport_data_{airline.lower().replace(' ', '_')}.csv", | |
| mime='text/csv', | |
| ) | |
| output = BytesIO() | |
| with pd.ExcelWriter(output, engine='openpyxl') as writer: | |
| df.to_excel(writer, index=False, sheet_name='PassportData') | |
| col2.download_button( | |
| label="Download data as Excel", | |
| data=output.getvalue(), | |
| file_name=f"passport_data_{airline.lower().replace(' ', '_')}.xlsx", | |
| mime='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', | |
| ) | |
| if __name__ == "__main__": | |
| main() | |