CoolFace
Apppublic

imran-decoder/filecrackhead1

sourceHugging Faceupdated 6mo agoView on Hugging Face
0likes
ocr_converter.py196 linesDownload Raw Back to services
1"""2OCR Converter: PDF → DOCX with full OCR pipeline.3Premium feature: uses pdf2image + pytesseract for scanned/image PDFs.4 5Two modes are available via the caller:6  use_ocr=False → Fast text-layer extraction (PyMuPDF) — free tier7  use_ocr=True  → Full OCR pipeline — premium/billable tier8 9This module handles ONLY the OCR path. The non-OCR path is in pdf_converter.py.10"""11 12import io13from pathlib import Path14from typing import Optional15 16import fitz  # PyMuPDF17import pytesseract18from docx import Document19from docx.shared import Pt20from PIL import Image21 22from app.config import get_settings23from app.utils.logging_utils import get_logger24 25logger = get_logger(__name__)26settings = get_settings()27 28 29def _has_extractable_text(pdf_path: Path, sample_pages: int = 3) -> bool:30    """31    Heuristic: check if the first N pages have meaningful embedded text.32    Returns True if text extraction is viable (no OCR needed).33    """34    total_chars = 035    with fitz.open(str(pdf_path)) as doc:36        for i, page in enumerate(doc):37            if i >= sample_pages:38                break39            total_chars += len(page.get_text().strip())40        page_count = doc.page_count41    # If average > 50 chars per sampled page, assume text-layer PDF42    return total_chars / max(sample_pages, page_count) > 5043 44 45def _extract_text_layer(pdf_path: Path) -> list[str]:46    """Fast: extract text using PyMuPDF's embedded text layer."""47    with fitz.open(str(pdf_path)) as doc:48        pages = [page.get_text() for page in doc]49    return pages50 51 52def _ocr_pages(53    pdf_path: Path,54    language: str = "eng",55    dpi: int = 300,56    timeout: Optional[int] = None,57) -> list[str]:58    """59    Premium OCR pipeline:60    1. Render each PDF page to a high-res PIL image (PyMuPDF)61    2. Run pytesseract on each image62    3. Return list of extracted text per page63    """64    pytesseract.pytesseract.tesseract_cmd = settings.TESSERACT_PATH65 66    logger.info(67        "Starting OCR pipeline",68        extra={"pdf": str(pdf_path), "dpi": dpi, "language": language},69    )70 71    zoom = dpi / 72.072    mat = fitz.Matrix(zoom, zoom)73    extracted = []74 75    with fitz.open(str(pdf_path)) as doc:76        if len(doc) == 0:77            raise RuntimeError("PDF returned no pages — cannot perform OCR.")78 79        for i, page in enumerate(doc):80            logger.debug("OCR processing page %d/%d", i + 1, len(doc))81            pix = page.get_pixmap(matrix=mat, alpha=False)82            83            # Convert PyMuPDF Pixmap to PIL Image84            if pix.n - pix.alpha < 4:       # this is RGB or gray85                mode = "RGB"86            else:                           # CMYK87                mode = "RGBA"88                89            img = Image.frombytes(mode, [pix.width, pix.height], pix.samples)90            if mode != "RGB":91                img = img.convert("RGB")92                93            # Pre-process: convert to grayscale for better OCR accuracy94            gray = img.convert("L")95            text = pytesseract.image_to_string(96                gray,97                lang=language,98                config="--psm 3 --oem 3",  # psm3=auto, oem3=LSTM99                timeout=timeout,100            )101            extracted.append(text)102 103    return extracted104 105 106def _build_docx_from_pages(pages: list[str], output_path: Path) -> Path:107    """108    Assemble a DOCX from a list of per-page text strings.109    Adds page separators between pages for readability.110    """111    doc = Document()112 113    # Apply a readable base font114    style = doc.styles["Normal"]115    style.font.name = "Calibri"116    style.font.size = Pt(11)117 118    for i, page_text in enumerate(pages):119        if i > 0:120            # Page break between pages121            doc.add_page_break()122 123        for line in page_text.splitlines():124            stripped = line.strip()125            if stripped:126                para = doc.add_paragraph(stripped)127                para.style = doc.styles["Normal"]128 129    doc.save(str(output_path))130    return output_path131 132 133def pdf_to_docx_no_ocr(input_path: Path, output_path: Path, **_) -> Path:134    """135    FREE TIER: Extract embedded text from PDF → DOCX using PyMuPDF.136    Fast (~seconds). Works only on text-layer PDFs. Scanned PDFs produce empty/garbled output.137    """138    pages = _extract_text_layer(input_path)139    return _build_docx_from_pages(pages, output_path)140 141 142def pdf_to_docx_with_ocr(143    input_path: Path,144    output_path: Path,145    language: str = "eng",146    dpi: int = 300,147    **_,148) -> Path:149    """150    PREMIUM TIER: Full OCR pipeline — renders pages to images then runs Tesseract.151    Accurate for scanned PDFs, handwriting, and image-based documents.152    Slower (~10-60s per page depending on DPI and content complexity).153 154    Parameters155    ----------156    input_path : Path to the source PDF157    output_path : Path for the output DOCX158    language : Tesseract language code(s), e.g. "eng", "eng+ara"159    dpi : Render DPI — higher = better accuracy, slower. Default 300 is optimal.160    """161    eff_language = language or settings.OCR_LANGUAGE162    eff_dpi = dpi or settings.OCR_DPI163 164    pages = _ocr_pages(165        input_path,166        language=eff_language,167        dpi=eff_dpi,168        timeout=settings.OCR_TIMEOUT_SECONDS,169    )170    return _build_docx_from_pages(pages, output_path)171 172 173def pdf_to_docx_smart(174    input_path: Path,175    output_path: Path,176    use_ocr: bool = False,177    language: str = "eng",178    dpi: int = 300,179    **kwargs,180) -> Path:181    """182    Smart dispatcher:183    - If use_ocr=False → structured extraction using pdf2docx (free)184    - If use_ocr=True  → full OCR pipeline (premium)185 186    This is the primary entry point used by the conversion engine.187    """188    if use_ocr:189        logger.info("Using OCR pipeline for PDF→DOCX conversion")190        return pdf_to_docx_with_ocr(input_path, output_path, language=language, dpi=dpi)191    else:192        logger.info("Using pdf2docx structured extraction for PDF→DOCX conversion")193        from app.services.pdf_converter import pdf_to_docx194        work_dir = kwargs.get("work_dir", input_path.parent)195        return pdf_to_docx(input_path, output_path, work_dir=work_dir)196