Khanishka/AI-document-analyzer-mini

1"""2extractor.py3------------4Handles all document extraction:5  - PDF  → text + tables (pdfplumber)6  - DOCX → text (python-docx)7  - Image → OCR text (pytesseract)8"""9 10import io11import time12import pdfplumber13from PIL import Image14from docx import Document15 16 17# ─────────────────────────────────────────────18# PDF EXTRACTION19# ─────────────────────────────────────────────20 21def extract_pdf(file_bytes: bytes) -> dict:22    """23    Extract text, tables, and metadata from a PDF file.24 25    Args:26        file_bytes: Raw bytes of the uploaded PDF.27 28    Returns:29        dict with keys: text, tables, page_count, processing_time_ms30    """31    start = time.time()32    full_text = []33    all_tables = []34 35    with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:36        page_count = len(pdf.pages)37 38        for page_num, page in enumerate(pdf.pages, start=1):39            # Extract plain text40            page_text = page.extract_text()41            if page_text:42                full_text.append(page_text.strip())43 44            # Extract tables → list of list of rows45            tables = page.extract_tables()46            for table in tables:47                if table:48                    # Filter out completely empty tables49                    non_empty_rows = []50                    for row in table:51                        cleaned_row = [cell.strip() if cell else "" for cell in row]52                        # Only keep rows that have at least one non-empty cell53                        if any(c for c in cleaned_row):54                            non_empty_rows.append(cleaned_row)55                    if non_empty_rows:56                        all_tables.append({57                            "page": page_num,58                            "data": non_empty_rows59                        })60 61    elapsed_ms = round((time.time() - start) * 1000, 2)62 63    return {64        "text": "\n\n".join(full_text),65        "tables": all_tables,66        "page_count": page_count,67        "processing_time_ms": elapsed_ms68    }69 70 71# ─────────────────────────────────────────────72# DOCX EXTRACTION73# ─────────────────────────────────────────────74 75def extract_docx(file_bytes: bytes) -> dict:76    """77    Extract text and paragraph count from a DOCX file.78 79    Args:80        file_bytes: Raw bytes of the uploaded DOCX.81 82    Returns:83        dict with keys: text, paragraph_count, processing_time_ms84    """85    start = time.time()86    doc = Document(io.BytesIO(file_bytes))87 88    paragraphs = []89    for para in doc.paragraphs:90        stripped = para.text.strip()91        if stripped:92            paragraphs.append(stripped)93 94    # Also extract text from tables inside DOCX95    table_texts = []96    for table in doc.tables:97        for row in table.rows:98            row_text = " | ".join(99                cell.text.strip() for cell in row.cells if cell.text.strip()100            )101            if row_text:102                table_texts.append(row_text)103 104    all_text_parts = paragraphs + table_texts105    elapsed_ms = round((time.time() - start) * 1000, 2)106 107    return {108        "text": "\n\n".join(all_text_parts),109        "paragraph_count": len(paragraphs),110        "processing_time_ms": elapsed_ms111    }112 113 114# ─────────────────────────────────────────────115# IMAGE OCR EXTRACTION116# ─────────────────────────────────────────────117 118def extract_image(file_bytes: bytes) -> dict:119    """Image OCR not available in mini version."""120    return {121        "text": "",122        "error": "Image OCR not supported in this deployment. Use PDF or DOCX.",123        "processing_time_ms": 0124    }125 126 127def extract(file_bytes: bytes, file_type: str) -> dict:128    """129    Route extraction based on file type.130 131    Args:132        file_bytes : Raw bytes of the file.133        file_type  : One of 'pdf', 'docx', 'image'134 135    Returns:136        Extraction result dict (varies by type).137 138    Raises:139        ValueError: If file_type is not supported.140    """141    file_type = file_type.lower().strip()142 143    if file_type == "pdf":144        return extract_pdf(file_bytes)145    elif file_type in ("docx", "doc"):146        return extract_docx(file_bytes)147    elif file_type in ("image", "png", "jpg", "jpeg", "bmp", "tiff", "webp"):148        return extract_image(file_bytes)149    else:150        raise ValueError(151            f"Unsupported file type: '{file_type}'. "152            "Supported types: pdf, docx, image (png/jpg/jpeg/bmp/tiff/webp)"153        )154