Spaces:
Sleeping
Sleeping
| 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, | |
| verbose=False | |
| ) | |
| logger.debug("EasyOCR initialized.") | |
| # --------------------------------------------------- | |
| # VISUAL GIVEN NAME EXTRACTION | |
| # --------------------------------------------------- | |
| 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 "" | |
| candidate = re.sub(r'[^A-Za-z\s]', '', candidate) | |
| candidate = re.sub(r'\s+', ' ', candidate).strip() | |
| candidate = re.sub(r'([A-Z]{2,})[K]$', r'\1', candidate) | |
| return candidate.strip() | |
| return "" | |
| except Exception as e: | |
| logger.error(f"Given Names extraction failed: {e}") | |
| return "" | |
| # --------------------------------------------------- | |
| # MRZ EXTRACTION (SMART FALLBACK) | |
| # --------------------------------------------------- | |
| def extract_mrz_easyocr(self, img_path): | |
| """Find MRZ lines using EasyOCR when Tesseract is missing.""" | |
| try: | |
| # Use EasyOCR to read ALL text | |
| allow = st.ascii_uppercase + st.digits + "<" | |
| results = self.reader.readtext(img_path, detail=0, allowlist=allow) | |
| # 1. Clean and filter candidates | |
| candidates = [] | |
| for res in results: | |
| # Remove spaces and keep uppercase/digits/< | |
| clean = "".join([c for c in res.upper().replace(" ", "") if c in allow]) | |
| if len(clean) >= 30: | |
| candidates.append(clean) | |
| # 2. Look for Line 1 (starts with P<) | |
| line1, line2 = None, None | |
| for i, c in enumerate(candidates): | |
| if c.startswith("P<"): | |
| line1 = c | |
| # Line 2 is almost always the next one or one after | |
| if i + 1 < len(candidates): | |
| line2 = candidates[i+1] | |
| break | |
| # 3. If no P< found, look for any two 44-char lines | |
| if not line1 and len(candidates) >= 2: | |
| for i in range(len(candidates)-1): | |
| if len(candidates[i]) >= 40 and len(candidates[i+1]) >= 40: | |
| line1, line2 = candidates[i], candidates[i+1] | |
| break | |
| return line1, line2 | |
| except Exception as e: | |
| logger.error(f"EasyOCR fallback error: {e}") | |
| return None, None | |
| def extract_mrz_from_roi(self, img_path): | |
| try: | |
| # First attempt: PassportEye (If Tesseract works) | |
| try: | |
| mrz = read_mrz(img_path, save_roi=True) | |
| if mrz and getattr(mrz, 'valid_score', 0) > 20: | |
| l1 = clean_mrz_line(getattr(mrz, 'line1', '')) | |
| l2 = clean_mrz_line(getattr(mrz, 'line2', '')) | |
| if l1 and l2: | |
| return l1, l2, mrz | |
| except: | |
| pass | |
| # Second attempt: EasyOCR Manual Search | |
| l1, l2 = self.extract_mrz_easyocr(img_path) | |
| if l1 and l2: | |
| # Standardize to 44 chars for FallbackMRZ | |
| l1_clean = clean_mrz_line(l1) | |
| l2_clean = clean_mrz_line(l2) | |
| return l1_clean, l2_clean, FallbackMRZ(l1_clean, l2_clean) | |
| return None, None, None | |
| 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): | |
| return None | |
| line1, line2, mrz = self.extract_mrz_from_roi(img_path) | |
| if mrz is 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 | |
| } | |
| data = { | |
| "surname": clean_name_field(getattr(mrz, "surname", "")), | |
| "name": self.extract_given_names_from_visual(img_path) or clean_name_field(getattr(mrz, "names", getattr(mrz, "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 | |
| def process_pdf(self, pdf_path, progress_callback=None, airline="iraqi"): | |
| os.makedirs(TEMP_DIR, exist_ok=True) | |
| try: | |
| from pdf2image import pdfinfo_from_path, convert_from_path | |
| info = pdfinfo_from_path(pdf_path) | |
| results = [] | |
| for page in range(1, info["Pages"] + 1): | |
| if progress_callback: progress_callback(page / info["Pages"]) | |
| imgs = convert_from_path(pdf_path, dpi=150, first_page=page, last_page=page) | |
| if not imgs: continue | |
| tmp = os.path.join(TEMP_DIR, f"tmp_{page}.jpg") | |
| imgs[0].save(tmp, "JPEG", quality=70) | |
| res = self.get_data(tmp, airline=airline) | |
| if res: | |
| res["page_number"] = page | |
| results.append(res) | |
| if os.path.exists(tmp): os.remove(tmp) | |
| return results | |
| except Exception as e: | |
| logger.error(f"PDF failed: {e}"); return [] | |