core/extractor.py

import io
import logging
from pathlib import Path

logger = logging.getLogger(__name__)


def extract_with_tesseract(file_bytes: bytes, file_extension: str) -> dict:
    """Extract text using Tesseract OCR — Arabic + English."""
    try:
        from PIL import Image
        import pytesseract

        extracted = {
            "source": "tesseract",
            "pages": [],
            "full_text": "",
            "tables": [],
            "key_value_pairs": [],
            "language": "ar",
            "confidence": 0.0,
            "word_count": 0
        }

        if file_extension.lower() == ".pdf":
            try:
                from pdf2image import convert_from_bytes
                images = convert_from_bytes(file_bytes, dpi=150)
            except Exception as e:
                logger.warning(f"PDF conversion failed: {e}")
                extracted["full_text"] = (
                    "[PDF extraction requires poppler. "
                    "Please upload a PNG or JPG image instead.]"
                )
                extracted["pages"].append({
                    "page_number": 1,
                    "lines": [],
                    "words": [],
                })
                return extracted
        else:
            image = Image.open(io.BytesIO(file_bytes))
            # Convert to RGB if needed
            if image.mode not in ("RGB", "L"):
                image = image.convert("RGB")
            images = [image]

        full_text_parts = []

        for page_idx, image in enumerate(images):
            # Try Arabic+English, fall back to English only
            try:
                text = pytesseract.image_to_string(
                    image, lang="ara+eng", config="--psm 3"
                )
            except Exception:
                try:
                    text = pytesseract.image_to_string(
                        image, lang="eng", config="--psm 3"
                    )
                except Exception:
                    text = pytesseract.image_to_string(image, config="--psm 3")

            lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
            full_text_parts.extend(lines)

            page_data = {
                "page_number": page_idx + 1,
                "lines": [{"text": ln, "confidence": 0.75} for ln in lines],
                "words": [
                    {"text": w, "confidence": 0.75}
                    for ln in lines
                    for w in ln.split()
                ],
            }
            extracted["pages"].append(page_data)

        extracted["full_text"] = "\n".join(full_text_parts)
        extracted["word_count"] = len(extracted["full_text"].split())
        extracted["confidence"] = 0.75 if extracted["word_count"] > 0 else 0.0

        return extracted

    except Exception as e:
        logger.error(f"Tesseract extraction failed: {e}")
        raise


def extract_document(
    file_bytes: bytes, file_extension: str, filename: str = ""
) -> dict:
    """Main extraction entry point — always uses Tesseract."""
    logger.info(f"Extracting: {filename}")

    result = extract_with_tesseract(file_bytes, file_extension)
    result["filename"] = filename
    result["file_size_bytes"] = len(file_bytes)
    result["file_extension"] = file_extension
    return result