first commit
This commit is contained in:
@@ -0,0 +1,110 @@
|
||||
"""Marker/OCR integration for document processing."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
import io
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
import httpx
|
||||
|
||||
from app.config import get_settings
|
||||
import app.db.database as _database_mod
|
||||
|
||||
|
||||
class MarkerProcessor:
|
||||
"""Handles OCR via Marker API with polygon output."""
|
||||
|
||||
def __init__(self):
|
||||
self.settings = get_settings()
|
||||
|
||||
async def process_pdf(self, file_data: bytes, filename: str) -> dict[str, Any]:
|
||||
"""Send PDF to Marker API for OCR with polygon extraction."""
|
||||
async with httpx.AsyncClient(timeout=300) as client:
|
||||
resp = await client.post(
|
||||
f"{self.settings.marker_api_url}/convert",
|
||||
files={"file": (filename, file_data, "application/pdf")},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
async def process_pdf_url(self, pdf_url: str, filename: str) -> dict[str, Any]:
|
||||
"""Process PDF from URL via Marker API."""
|
||||
async with httpx.AsyncClient(timeout=300) as client:
|
||||
resp = await client.post(
|
||||
f"{self.settings.marker_api_url}/convert",
|
||||
json={"pdf_url": pdf_url},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
async def parse_marker_json(self, marker_output: dict, doc_id: str) -> list[dict]:
|
||||
"""Parse Marker JSON output into vectorized chunks with polygon data."""
|
||||
pages = marker_output.get("pages", [])
|
||||
chunks = []
|
||||
|
||||
for page in pages:
|
||||
page_num = page.get("meta", {}).get("page_num", page.get("page", 0))
|
||||
text_lines = page.get("text_lines", [])
|
||||
|
||||
for block_idx, tl in enumerate(text_lines):
|
||||
content = tl.get("text", "")
|
||||
polygon = tl.get("bbox") or tl.get("polygon")
|
||||
block_type = tl.get("type", "text")
|
||||
|
||||
if not content or not isinstance(content, str) or not content.strip():
|
||||
continue
|
||||
|
||||
chunks.append({
|
||||
"content": content.strip(),
|
||||
"page_num": page_num,
|
||||
"block_index": block_idx,
|
||||
"polygon": polygon,
|
||||
"chunk_type": block_type,
|
||||
})
|
||||
|
||||
if chunks:
|
||||
await _database_mod.db.batch_chunk(doc_id, chunks)
|
||||
|
||||
return chunks
|
||||
|
||||
async def process_document_file(
|
||||
self, file_data: bytes, filename: str, doc_id: str
|
||||
) -> dict[str, Any]:
|
||||
"""Process document file and return structured result."""
|
||||
ext = Path(filename).suffix.lower()
|
||||
|
||||
if ext == ".pdf":
|
||||
result = await self.process_pdf(file_data, filename)
|
||||
pages_data = result.get("pages", [])
|
||||
|
||||
chunks = await self.parse_marker_json({"pages": pages_data}, doc_id)
|
||||
return {
|
||||
"success": result.get("success", True),
|
||||
"doc_id": doc_id,
|
||||
"page_count": result.get("page_count", len(pages_data)),
|
||||
"chunks": len(chunks),
|
||||
"ocr_model": result.get("ocr_model", "deepseek-ocr"),
|
||||
}
|
||||
|
||||
elif ext in (".txt", ".pdf", ".md"):
|
||||
text = file_data.decode("utf-8", errors="replace")
|
||||
paragraphs = re.split(r'\n\s*\n', text)
|
||||
chunks = []
|
||||
for i, para in enumerate(paragraphs):
|
||||
if len(para.strip()) > 20:
|
||||
chunks.append({
|
||||
"content": para.strip(),
|
||||
"page_num": 0,
|
||||
"block_index": i,
|
||||
"polygon": None,
|
||||
"chunk_type": "text",
|
||||
})
|
||||
|
||||
if chunks:
|
||||
await _database_mod.db.batch_chunk(doc_id, chunks)
|
||||
|
||||
return {"success": True, "doc_id": doc_id, "chunks": len(chunks), "page_count": 1}
|
||||
|
||||
return {"success": False, "error": f"Unsupported file type: {ext}"}
|
||||
Reference in New Issue
Block a user