107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import base64
|
||
|
|
import io
|
||
|
|
import os
|
||
|
|
import tempfile
|
||
|
|
import subprocess
|
||
|
|
import requests
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
OLLAMA_HOST = os.environ.get("DEEPSEEK_OLLAMA_HOST", "http://localhost:11434")
|
||
|
|
DEEPSEEK_MODEL = os.environ.get("DEEPSEEK_OCR_MODEL", "deepseek-ocr:latest")
|
||
|
|
DEEPSEEK_PROMPT = os.environ.get("DEEPSEEK_OCR_PROMPT", "<|grounding|>Free OCR.")
|
||
|
|
|
||
|
|
|
||
|
|
def _page_to_image(raw_bytes: bytes, page_num: int = 0) -> bytes | None:
|
||
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
||
|
|
result = subprocess.run(
|
||
|
|
[
|
||
|
|
"pdftoppm", "-png",
|
||
|
|
"-f", str(page_num + 1),
|
||
|
|
"-l", str(page_num + 1),
|
||
|
|
"-",
|
||
|
|
os.path.join(tmpdir, "page"),
|
||
|
|
],
|
||
|
|
input=raw_bytes,
|
||
|
|
capture_output=True,
|
||
|
|
timeout=120,
|
||
|
|
)
|
||
|
|
if result.returncode != 0:
|
||
|
|
return None
|
||
|
|
out_files = os.listdir(tmpdir)
|
||
|
|
if not out_files:
|
||
|
|
return None
|
||
|
|
img_path = os.path.join(tmpdir, out_files[0])
|
||
|
|
img = Image.open(img_path)
|
||
|
|
buf = io.BytesIO()
|
||
|
|
img.save(buf, format="PNG")
|
||
|
|
return buf.getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
def _image_to_text(image_bytes: bytes) -> str:
|
||
|
|
b64 = base64.b64encode(image_bytes).decode()
|
||
|
|
resp = requests.post(
|
||
|
|
f"{OLLAMA_HOST}/api/generate",
|
||
|
|
json={
|
||
|
|
"model": DEEPSEEK_MODEL,
|
||
|
|
"prompt": DEEPSEEK_PROMPT,
|
||
|
|
"images": [b64],
|
||
|
|
"stream": False,
|
||
|
|
"options": {"temperature": 0},
|
||
|
|
},
|
||
|
|
timeout=120,
|
||
|
|
)
|
||
|
|
resp.raise_for_status()
|
||
|
|
return resp.json().get("response", "")
|
||
|
|
|
||
|
|
|
||
|
|
def ocr_pdf(raw_bytes: bytes, max_pages: int | None = None, page_range: str | None = None) -> dict:
|
||
|
|
result = {
|
||
|
|
"success": True,
|
||
|
|
"text": "",
|
||
|
|
"pages": [],
|
||
|
|
"page_count": 0,
|
||
|
|
}
|
||
|
|
pages_to_process = _parse_page_range(page_range, max_pages, raw_bytes)
|
||
|
|
for page_num in pages_to_process:
|
||
|
|
img_bytes = _page_to_image(raw_bytes, page_num)
|
||
|
|
if img_bytes is None:
|
||
|
|
continue
|
||
|
|
text = _image_to_text(img_bytes)
|
||
|
|
result["pages"].append({"page": page_num, "text": text})
|
||
|
|
result["text"] += f"\n\n--- PAGE {page_num + 1} ---\n\n{text}"
|
||
|
|
result["page_count"] = len(result["pages"])
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _parse_page_range(page_range: str | None, max_pages: int | None, raw_bytes: bytes) -> list[int]:
|
||
|
|
import tempfile, os, subprocess
|
||
|
|
total = 0
|
||
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as f:
|
||
|
|
f.write(raw_bytes)
|
||
|
|
f.flush()
|
||
|
|
r = subprocess.run(
|
||
|
|
["pdfinfo", f.name],
|
||
|
|
capture_output=True, text=True, timeout=30,
|
||
|
|
)
|
||
|
|
os.unlink(f.name)
|
||
|
|
for line in r.stdout.splitlines():
|
||
|
|
if line.startswith("Pages:"):
|
||
|
|
total = int(line.split(":")[1].strip())
|
||
|
|
break
|
||
|
|
if page_range:
|
||
|
|
pages = set()
|
||
|
|
for part in page_range.split(","):
|
||
|
|
part = part.strip()
|
||
|
|
if "-" in part:
|
||
|
|
a, b = part.split("-", 1)
|
||
|
|
for p in range(int(a.strip()), int(b.strip()) + 1):
|
||
|
|
pages.add(p)
|
||
|
|
else:
|
||
|
|
pages.add(int(part))
|
||
|
|
return sorted(pages)
|
||
|
|
if max_pages is not None:
|
||
|
|
return list(range(min(max_pages, total)))
|
||
|
|
return list(range(total))
|