Spaces:
Sleeping
Sleeping
| """ | |
| Text Extraction Module | |
| Handles PDF, DOCX, and Image file types with layout-aware extraction. | |
| """ | |
| import io | |
| import logging | |
| from pathlib import Path | |
| from typing import Optional | |
| import pdfplumber | |
| import fitz # PyMuPDF | |
| from docx import Document as DocxDocument | |
| from PIL import Image | |
| from app.modules.ocr_processor import OCRProcessor | |
| logger = logging.getLogger(__name__) | |
| class TextExtractor: | |
| """ | |
| Routes documents to the appropriate extraction pipeline based on file type. | |
| Supports PDF (text-based and scanned), DOCX, PNG, and JPG. | |
| """ | |
| SUPPORTED_TYPES = { | |
| "application/pdf": "_extract_pdf", | |
| "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "_extract_docx", | |
| "image/png": "_extract_image", | |
| "image/jpeg": "_extract_image", | |
| "image/jpg": "_extract_image", | |
| } | |
| EXTENSION_MAP = { | |
| ".pdf": "application/pdf", | |
| ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", | |
| ".png": "image/png", | |
| ".jpg": "image/jpeg", | |
| ".jpeg": "image/jpeg", | |
| } | |
| def __init__(self): | |
| self.ocr = OCRProcessor() | |
| def detect_mime_type(self, filename: str, content: bytes) -> str: | |
| """Detect MIME type from filename extension with fallback to content sniffing.""" | |
| ext = Path(filename).suffix.lower() | |
| if ext in self.EXTENSION_MAP: | |
| return self.EXTENSION_MAP[ext] | |
| # Content-based sniffing | |
| if content[:4] == b"%PDF": | |
| return "application/pdf" | |
| if content[:2] in (b"PK",): | |
| return "application/vnd.openxmlformats-officedocument.wordprocessingml.document" | |
| if content[:8] in (b"\x89PNG\r\n\x1a\n",): | |
| return "image/png" | |
| if content[:3] == b"\xff\xd8\xff": | |
| return "image/jpeg" | |
| raise ValueError(f"Unsupported or undetectable file type for: {filename}") | |
| def extract(self, content: bytes, filename: str) -> str: | |
| """ | |
| Main entry point. Detects file type and routes to correct extractor. | |
| Returns cleaned, normalized text string. | |
| """ | |
| mime = self.detect_mime_type(filename, content) | |
| method_name = self.SUPPORTED_TYPES.get(mime) | |
| if not method_name: | |
| raise ValueError(f"Unsupported MIME type: {mime}") | |
| method = getattr(self, method_name) | |
| logger.info(f"Extracting text from {filename} (MIME: {mime})") | |
| return method(content, filename) | |
| # ─── PDF Extraction ──────────────────────────────────────────────────────── | |
| def _extract_pdf(self, content: bytes, filename: str) -> str: | |
| """ | |
| Extract text from PDF. Tries selectable text first; falls back to OCR. | |
| Uses pdfplumber for layout-aware extraction. | |
| """ | |
| text = self._extract_pdf_selectable(content) | |
| if not text or len(text.strip()) < 50: | |
| logger.info(f"PDF '{filename}' appears scanned. Falling back to OCR.") | |
| text = self._extract_pdf_ocr(content) | |
| return text | |
| def _extract_pdf_selectable(self, content: bytes) -> str: | |
| """Extract selectable text from PDF using pdfplumber (layout-aware).""" | |
| pages_text = [] | |
| try: | |
| with pdfplumber.open(io.BytesIO(content)) as pdf: | |
| for page in pdf.pages: | |
| page_text = page.extract_text(x_tolerance=3, y_tolerance=3) | |
| if page_text: | |
| pages_text.append(page_text.strip()) | |
| except Exception as e: | |
| logger.warning(f"pdfplumber extraction failed: {e}") | |
| return "\n\n".join(pages_text) | |
| def _extract_pdf_ocr(self, content: bytes) -> str: | |
| """ | |
| Render each PDF page as an image and apply OCR. | |
| Uses PyMuPDF (fitz) for reliable PDF→image rendering. | |
| """ | |
| pages_text = [] | |
| try: | |
| doc = fitz.open(stream=content, filetype="pdf") | |
| for page_num in range(len(doc)): | |
| page = doc[page_num] | |
| # Render at 300 DPI for good OCR quality | |
| mat = fitz.Matrix(300 / 72, 300 / 72) | |
| pix = page.get_pixmap(matrix=mat, alpha=False) | |
| img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) | |
| page_text = self.ocr.extract(img) | |
| if page_text.strip(): | |
| pages_text.append(page_text.strip()) | |
| doc.close() | |
| except Exception as e: | |
| logger.error(f"PDF OCR failed: {e}") | |
| return "\n\n".join(pages_text) | |
| # ─── DOCX Extraction ─────────────────────────────────────────────────────── | |
| def _extract_docx(self, content: bytes, filename: str) -> str: | |
| """ | |
| Extract text from DOCX preserving paragraph integrity. | |
| Maintains logical structure: headings are separated by blank lines. | |
| """ | |
| try: | |
| doc = DocxDocument(io.BytesIO(content)) | |
| paragraphs = [] | |
| for para in doc.paragraphs: | |
| text = para.text.strip() | |
| if text: | |
| # Preserve heading structure | |
| if para.style.name.startswith("Heading"): | |
| paragraphs.append(f"\n{text}\n") | |
| else: | |
| paragraphs.append(text) | |
| # Also extract from tables | |
| for table in doc.tables: | |
| for row in table.rows: | |
| row_text = " | ".join( | |
| cell.text.strip() for cell in row.cells if cell.text.strip() | |
| ) | |
| if row_text: | |
| paragraphs.append(row_text) | |
| return "\n".join(paragraphs) | |
| except Exception as e: | |
| logger.error(f"DOCX extraction failed: {e}") | |
| return "" | |
| # ─── Image Extraction ────────────────────────────────────────────────────── | |
| def _extract_image(self, content: bytes, filename: str) -> str: | |
| """Extract text from images using OCR with preprocessing.""" | |
| try: | |
| img = Image.open(io.BytesIO(content)) | |
| return self.ocr.extract(img) | |
| except Exception as e: | |
| logger.error(f"Image extraction failed: {e}") | |
| return "" | |