from __future__ import annotations import uuid from pathlib import Path import cv2 import fitz # PyMuPDF import numpy as np from sqlalchemy.orm import Session from app.core.config import settings from app.core.logging_config import get_logger from app.models.document import Document, DocumentImage, DocumentPage, DocumentTable, DocumentTextBlock from app.repositories.document_repository import ( DocumentImageRepository, DocumentPageRepository, DocumentRepository, DocumentTableRepository, DocumentTextBlockRepository, ) from app.storage.provider import get_storage_provider logger = get_logger(__name__) class NativePDFService: """Extract text, fonts, images, and layout from native (non-scanned) PDFs using PyMuPDF.""" def __init__(self, db: Session) -> None: self.db = db self.storage = get_storage_provider() self.doc_repo = DocumentRepository(db) self.page_repo = DocumentPageRepository(db) self.text_block_repo = DocumentTextBlockRepository(db) self.image_repo = DocumentImageRepository(db) self.table_repo = DocumentTableRepository(db) def process_pdf(self, document: Document) -> Document: """Process a native PDF document, extracting all content.""" file_path = self.storage.get_absolute_path(document.storage_path) try: pdf_doc = fitz.open(file_path) except Exception as e: logger.error("pdf_open_failed", document_id=str(document.id), error=str(e)) raise document.page_count = len(pdf_doc) document.is_scanned = self._is_scanned_pdf(pdf_doc) for page_num in range(len(pdf_doc)): page = pdf_doc[page_num] self._process_page(document, page, page_num + 1) pdf_doc.close() return document def _is_scanned_pdf(self, pdf_doc: fitz.Document) -> bool: """Determine if a PDF is scanned (image-based) or native.""" total_text_chars = 0 total_images = 0 for page_num in range(min(len(pdf_doc), 3)): page = pdf_doc[page_num] text = page.get_text("text") total_text_chars += len(text.strip()) total_images += len(page.get_images(full=True)) # If very little text but has images, likely scanned if total_text_chars < 50 and total_images > 0: return True return False def _process_page(self, document: Document, page: fitz.Page, page_number: int) -> DocumentPage: """Process a single PDF page.""" rect = page.rect width = rect.width height = rect.height # Save page as image for reference pix = page.get_pixmap(dpi=150) image_filename = f"{document.id}_page_{page_number}.png" image_data = pix.tobytes("png") image_path = self.storage.save_file(image_data, "images", image_filename) # Get full text content text_content = page.get_text("text") doc_page = self.page_repo.create_page( document_id=document.id, page_number=page_number, width=width, height=height, image_path=image_path, text_content=text_content, ) # Extract text blocks with font information self._extract_text_blocks(doc_page, page) # Extract images self._extract_images(doc_page, page, document) # Detect headers and footers self._detect_headers_footers(doc_page, page) return doc_page def _extract_text_blocks(self, doc_page: DocumentPage, page: fitz.Page) -> None: """Extract text blocks with positioning and font information.""" blocks = page.get_text("rawdict")["blocks"] sequence = 0 chunks = [] for block in blocks: if block["type"] != 0: continue current_chunk = None space_count = 0 for line in block.get("lines", []): for span in line.get("spans", []): font_size = span.get("size", 12.0) # Ignore massive text (like diagonal watermarks) if font_size > 60: continue space_threshold = font_size * 1.5 font_family = span.get("font", None) color_int = span.get("color", 0) font_color = f"#{color_int:06x}" if isinstance(color_int, int) else None flags = span.get("flags", 0) styles = [] if flags & 1: styles.append("superscript") if flags & 2: styles.append("italic") if flags & 4: styles.append("serif") if flags & 8: styles.append("monospace") if flags & 16: styles.append("bold") font_style = ",".join(styles) if styles else "regular" font_info = { "family": font_family, "size": font_size, "color": font_color, "style": font_style } for char in span.get("chars", []): c = char["c"] bbox = char["bbox"] if c == ' ': space_count += 1 if space_count >= 2: if current_chunk and current_chunk["text"].strip(): chunks.append(current_chunk) current_chunk = None elif current_chunk: current_chunk["text"] += c current_chunk["bbox"][2] = max(current_chunk["bbox"][2], bbox[2]) current_chunk["bbox"][3] = max(current_chunk["bbox"][3], bbox[3]) continue else: space_count = 0 if current_chunk is None: current_chunk = {"text": c, "bbox": list(bbox), "font_info": font_info} continue prev_x1 = current_chunk["bbox"][2] distance = bbox[0] - prev_x1 if distance > space_threshold: if current_chunk["text"].strip(): chunks.append(current_chunk) current_chunk = {"text": c, "bbox": list(bbox), "font_info": font_info} else: current_chunk["text"] += c current_chunk["bbox"][2] = max(current_chunk["bbox"][2], bbox[2]) current_chunk["bbox"][3] = max(current_chunk["bbox"][3], bbox[3]) current_chunk["bbox"][1] = min(current_chunk["bbox"][1], bbox[1]) current_chunk["bbox"][0] = min(current_chunk["bbox"][0], bbox[0]) if current_chunk and current_chunk["text"].strip(): chunks.append(current_chunk) current_chunk = None for chunk in chunks: bbox = chunk["bbox"] font_info = chunk["font_info"] self.text_block_repo.create_text_block( page_id=doc_page.id, text=chunk["text"].strip(), x=bbox[0], y=bbox[1], width=bbox[2] - bbox[0], height=bbox[3] - bbox[1], font_family=font_info["family"], font_size=font_info["size"], font_color=font_info["color"], font_style=font_info["style"], block_type="text", sequence=sequence, ) sequence += 1 def _extract_images(self, doc_page: DocumentPage, page: fitz.Page, document: Document) -> None: """Extract embedded images from a PDF page.""" image_list = page.get_images(full=True) for img_index, img_info in enumerate(image_list): xref = img_info[0] try: base_image = page.parent.extract_image(xref) if not base_image: continue image_bytes = base_image["image"] ext = base_image.get("ext", "png") img_filename = f"{document.id}_page_{doc_page.page_number}_img_{img_index}.{ext}" img_storage_path = self.storage.save_file(image_bytes, "images", img_filename) # Try to get image position on page img_rects = page.get_image_rects(xref) if img_rects: rect = img_rects[0] x, y, x1, y1 = rect.x0, rect.y0, rect.x1, rect.y1 else: x, y, x1, y1 = 0, 0, base_image.get("width", 100), base_image.get("height", 100) # Determine image type based on position page_height = doc_page.height page_width = doc_page.width image_type = self._classify_image_type(x, y, x1, y1, page_width, page_height) self.image_repo.create_image( page_id=doc_page.id, x=x, y=y, width=x1 - x, height=y1 - y, image_path=img_storage_path, image_type=image_type, ) except Exception as e: logger.warning( "image_extraction_failed", page_id=str(doc_page.id), img_index=img_index, error=str(e), ) def _classify_image_type( self, x: float, y: float, x1: float, y1: float, page_width: float, page_height: float, ) -> str: """Classify an image as logo, figure, background, or stamp based on position and size.""" width = x1 - x height = y1 - y area_ratio = (width * height) / (page_width * page_height) if page_width > 0 and page_height > 0 else 0 # Background: covers most of the page if area_ratio > 0.8: return "background" # Logo: small image in top portion if y < page_height * 0.15 and area_ratio < 0.05: return "logo" # Stamp: small image in bottom-right if x > page_width * 0.6 and y > page_height * 0.7 and area_ratio < 0.05: return "stamp" return "figure" def _detect_headers_footers(self, doc_page: DocumentPage, page: fitz.Page) -> None: """Detect header and footer regions based on vertical position.""" page_height = page.rect.height header_threshold = page_height * 0.1 footer_threshold = page_height * 0.9 blocks = page.get_text("dict")["blocks"] header_seq = 0 footer_seq = 0 for block in blocks: if block["type"] != 0: continue bbox = block["bbox"] block_y = bbox[1] text_parts = [] for line in block.get("lines", []): for span in line.get("spans", []): t = span.get("text", "").strip() if t: text_parts.append(t) full_text = " ".join(text_parts) if not full_text.strip(): continue if block_y < header_threshold: self.text_block_repo.create_text_block( page_id=doc_page.id, text=full_text, x=bbox[0], y=bbox[1], width=bbox[2] - bbox[0], height=bbox[3] - bbox[1], block_type="header", sequence=header_seq, ) header_seq += 1 elif block_y > footer_threshold: self.text_block_repo.create_text_block( page_id=doc_page.id, text=full_text, x=bbox[0], y=bbox[1], width=bbox[2] - bbox[0], height=bbox[3] - bbox[1], block_type="footer", sequence=footer_seq, ) footer_seq += 1