import os import cv2 import numpy as np import easyocr import warnings import ssl import re from passporteye import read_mrz from pdf2image import convert_from_path from PIL import Image import string as st from src.utils import ( clean_string, clean_mrz_line, parse_date, get_country_name, get_sex, setup_logger, clean_name_field ) from src.fallback_mrz import FallbackMRZ from config.settings import USE_GPU, OCR_LANGUAGES, TEMP_DIR warnings.filterwarnings("ignore") logger = setup_logger(__name__) # Fix SSL issue (Mac EasyOCR model download fix) try: _create_unverified_https_context = ssl._create_unverified_context except AttributeError: pass else: ssl._create_default_https_context = _create_unverified_https_context class PassportExtractor: def __init__(self, use_gpu=USE_GPU, languages=None): self.languages = languages if languages else OCR_LANGUAGES base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) model_dir = os.path.join(base_dir, "data", "models") os.makedirs(model_dir, exist_ok=True) logger.debug(f"Initializing EasyOCR Reader (GPU={use_gpu})...") # Create custom user network directory to avoid permission issues user_network_dir = os.path.join(base_dir, "data", "easyocr_user_network") os.makedirs(user_network_dir, exist_ok=True) self.reader = easyocr.Reader( self.languages, gpu=use_gpu, model_storage_directory=model_dir, user_network_directory=user_network_dir ) logger.debug("EasyOCR initialized.") # --------------------------------------------------- # VISUAL GIVEN NAME EXTRACTION (PRIMARY SOURCE) # --------------------------------------------------- def extract_given_names_from_visual(self, img_path): try: results = self.reader.readtext(img_path, detail=0) lines = [r.strip() for r in results if r.strip()] for i, line in enumerate(lines): upper_line = line.upper() if "GIVEN" in upper_line and "NAME" in upper_line: if ":" in line: candidate = line.split(":")[1].strip() else: if i + 1 < len(lines): candidate = lines[i + 1].strip() else: return "" # Keep only letters and spaces, but preserve spaces between names candidate = re.sub(r'[^A-Za-z\s]', '', candidate) # Clean up extra spaces but keep single spaces between names candidate = re.sub(r'\s+', ' ', candidate).strip() # Remove trailing single letter only if it's clearly an OCR artifact (not part of a name) # This is more conservative - only removes single letters that are likely OCR errors candidate = re.sub(r'([A-Z]{2,})[K]$', r'\1', candidate) # K is common OCR error for < return candidate.strip() return "" except Exception as e: logger.error(f"Given Names extraction failed: {e}") return "" # --------------------------------------------------- # MRZ EXTRACTION # --------------------------------------------------- def extract_mrz_from_roi(self, img_path): try: mrz = read_mrz(img_path, save_roi=True) if not mrz: return None, None, None roi = mrz.aux["roi"] if roi.dtype != np.uint8: roi = (roi * 255).astype(np.uint8) img_resized = cv2.resize(roi, (1110, 140)) allow = st.ascii_uppercase + st.digits + "<" code = self.reader.readtext( img_resized, detail=0, allowlist=allow, batch_size=1, # Single image processing for speed workers=0, # Use main thread for stability decoder='greedy' # Faster decoding ) if len(code) < 2: return None, None, mrz line1 = clean_mrz_line(code[0]) line2 = clean_mrz_line(code[1]) return line1, line2, mrz except Exception as e: logger.error(f"MRZ extraction failed: {e}") return None, None, None # --------------------------------------------------- # MAIN DATA FUNCTION # --------------------------------------------------- def get_data(self, img_path, airline="iraqi"): if not os.path.exists(img_path): logger.error(f"File not found: {img_path}") return None line1, line2, mrz = self.extract_mrz_from_roi(img_path) if line1 and line2: mrz = FallbackMRZ(line1, line2) if mrz is None: logger.warning("MRZ not detected.") # Return data with placeholder values instead of None return { "surname": "•••", "name": "•••", "country": "•••", "nationality": "•••", "passport_number": "•••", "sex": "•••", "date_of_birth": "•••", "expiration_date": "•••", "personal_number": "•••", "mrz_full_string": "", "valid_score": 0, "mrz_found": False } surname = clean_name_field(getattr(mrz, "surname", "")) # 🔥 ALWAYS prefer visual name visual_name = self.extract_given_names_from_visual(img_path) if visual_name: name = visual_name else: name = clean_name_field( getattr(mrz, "names", getattr(mrz, "name", "")) ) # Final defensive cleanup - only remove trailing K which is a common OCR artifact name = re.sub(r'([A-Z]{2,})[K]$', r'\1', name) data = { "surname": surname, "name": name, "country": get_country_name(getattr(mrz, "country", "")), "nationality": get_country_name(getattr(mrz, "nationality", "")), "passport_number": clean_string(getattr(mrz, "number", "")), "sex": get_sex(getattr(mrz, "sex", "")), "date_of_birth": parse_date(getattr(mrz, "date_of_birth", ""), airline=airline), "expiration_date": parse_date(getattr(mrz, "expiration_date", ""), airline=airline), "mrz_full_string": (line1 or "") + (line2 or ""), "valid_score": getattr(mrz, "valid_score", 0), "mrz_found": True, } return data # --------------------------------------------------- # PDF PROCESSING # --------------------------------------------------- def process_pdf(self, pdf_path, progress_callback=None, airline="iraqi"): """ Memory-safe PDF processing for Streamlit free tier with fallback support. Converts PDF pages to images and extracts passport data from each page. Args: pdf_path (str): Path to the PDF file. progress_callback (function, optional): Progress callback function. airline (str): Airline format for date formatting ("iraqi", "default", "fly dubai", "fly baghdad"). Returns: list: List of dictionaries with extracted passport data per page. """ # Ensure temp directory exists before processing os.makedirs(TEMP_DIR, exist_ok=True) # Check file size for Streamlit free tier (max 10 MB) try: file_size = os.path.getsize(pdf_path) if file_size > 10 * 1024 * 1024: # 10 MB limit logger.error(f"PDF too large for free tier: {file_size / (1024*1024):.1f} MB") return [] except Exception as e: logger.error(f"Could not check file size: {e}") return [] results = [] try: # Try pdf2image first (primary method) from pdf2image import pdfinfo_from_path, convert_from_path info = pdfinfo_from_path(pdf_path) total_pages = info["Pages"] logger.debug(f"Processing PDF with {total_pages} pages (size: {file_size / (1024*1024):.1f} MB) using pdf2image") for page in range(1, total_pages + 1): try: # Update progress if callback provided if progress_callback: progress_callback(page / total_pages) # Convert ONE page at a time, optimize for speed images = convert_from_path( pdf_path, dpi=150, # Further reduced DPI for faster processing first_page=page, last_page=page, thread_count=1, # Single thread for stability use_pdftocairo=True # Faster backend ) image = images[0] # Save temporary image to TEMP_DIR with lower quality for speed temp_image_path = os.path.join(TEMP_DIR, f"temp_page_{page}.jpg") image.save(temp_image_path, "JPEG", quality=70, optimize=True) # Extract passport data with airline-specific formatting result = self.get_data(temp_image_path, airline=airline) if result: result["page_number"] = page results.append(result) logger.debug(f"Successfully extracted data from page {page}") # Delete temp image immediately to free memory if os.path.exists(temp_image_path): os.remove(temp_image_path) except Exception as e: logger.error(f"Error on page {page} with pdf2image: {e}") continue logger.debug(f"PDF processing finished with pdf2image. Valid pages: {len(results)}") except Exception as pdf2image_error: logger.warning(f"pdf2image failed: {pdf2image_error}. Trying fallback with PyMuPDF...") # Fallback to PyMuPDF (fitz) try: import fitz doc = fitz.open(pdf_path) total_pages = len(doc) logger.debug(f"Processing PDF with {total_pages} pages using PyMuPDF fallback") for i in range(total_pages): try: # Update progress if callback provided if progress_callback: progress_callback((i + 1) / total_pages) page = doc.load_page(i) pix = page.get_pixmap(dpi=200) # Convert to PIL Image img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) # Save temporary image to TEMP_DIR temp_image_path = os.path.join(TEMP_DIR, f"temp_page_{i+1}.png") img.save(temp_image_path, "PNG") # Extract passport data with airline-specific formatting result = self.get_data(temp_image_path, airline=airline) if result: result["page_number"] = i + 1 results.append(result) logger.debug(f"Successfully extracted data from page {i+1}") # Cleanup if os.path.exists(temp_image_path): os.remove(temp_image_path) except Exception as e: logger.error(f"Error on page {i+1} with PyMuPDF: {e}") continue doc.close() logger.debug(f"PDF processing finished with PyMuPDF fallback. Valid pages: {len(results)}") except Exception as fitz_error: logger.error(f"PyMuPDF fallback also failed: {fitz_error}") return [] return results