Changes committed
This commit is contained in:
292
docengine/app/services/pdf_service.py
Normal file
292
docengine/app/services/pdf_service.py
Normal file
@@ -0,0 +1,292 @@
|
||||
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("dict", flags=fitz.TEXT_PRESERVE_WHITESPACE)["blocks"]
|
||||
sequence = 0
|
||||
|
||||
for block in blocks:
|
||||
if block["type"] != 0: # Skip non-text blocks
|
||||
continue
|
||||
|
||||
block_text_parts = []
|
||||
font_info = {"family": None, "size": None, "color": None, "style": None}
|
||||
|
||||
for line in block.get("lines", []):
|
||||
for span in line.get("spans", []):
|
||||
text = span.get("text", "").strip()
|
||||
if text:
|
||||
block_text_parts.append(text)
|
||||
# Capture font info from the first non-empty span
|
||||
if font_info["family"] is None:
|
||||
font_info["family"] = span.get("font", None)
|
||||
font_info["size"] = span.get("size", None)
|
||||
color_int = span.get("color", 0)
|
||||
font_info["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_info["style"] = ",".join(styles) if styles else "regular"
|
||||
|
||||
full_text = " ".join(block_text_parts)
|
||||
if not full_text.strip():
|
||||
continue
|
||||
|
||||
bbox = block["bbox"]
|
||||
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],
|
||||
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
|
||||
Reference in New Issue
Block a user