src/extractor.py

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 []