333 lines
12 KiB
Python
333 lines
12 KiB
Python
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)
|
|
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
|