349 lines
14 KiB
Python
349 lines
14 KiB
Python
import io
|
|
import time
|
|
from typing import Any, Dict, List, Optional
|
|
from PIL import Image
|
|
|
|
try:
|
|
from rapidocr_onnxruntime import RapidOCR
|
|
_ocr_engine = RapidOCR()
|
|
_ocr_available = True
|
|
except Exception as e:
|
|
_ocr_engine = None
|
|
_ocr_available = False
|
|
print(f"Warning: RapidOCR not available: {e}")
|
|
|
|
|
|
def is_ocr_available() -> bool:
|
|
return _ocr_available and _ocr_engine is not None
|
|
|
|
|
|
def recognize_image_bytes(image_bytes: bytes) -> Dict[str, Any]:
|
|
if not is_ocr_available():
|
|
raise RuntimeError("RapidOCR engine is not available")
|
|
|
|
start_time = time.time()
|
|
img = Image.open(io.BytesIO(image_bytes))
|
|
width, height = img.size
|
|
|
|
# Run RapidOCR with parameters tuned for numbers, single digits, isolated symbols (- / .), and table cells
|
|
result, _ = _ocr_engine(
|
|
image_bytes,
|
|
box_thresh=0.08,
|
|
text_score=0.08,
|
|
unclip_ratio=2.5,
|
|
)
|
|
|
|
lines = []
|
|
if result:
|
|
for item in result:
|
|
# item format: [box, text, score]
|
|
# box: [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
|
|
box, text, score = item[0], item[1], float(item[2])
|
|
|
|
xs = [pt[0] for pt in box]
|
|
ys = [pt[1] for pt in box]
|
|
x_min, x_max = min(xs), max(xs)
|
|
y_min, y_max = min(ys), max(ys)
|
|
|
|
w = max(1.0, x_max - x_min)
|
|
h = max(1.0, y_max - y_min)
|
|
|
|
# Sub-split into words if text contains spaces
|
|
word_strings = text.split()
|
|
words = []
|
|
if len(word_strings) > 1:
|
|
char_w = w / max(1, len(text))
|
|
curr_x = x_min
|
|
for ws in word_strings:
|
|
w_width = len(ws) * char_w
|
|
words.append({
|
|
"text": ws,
|
|
"box": {"x": curr_x, "y": y_min, "width": w_width, "height": h},
|
|
"confidence": score,
|
|
})
|
|
curr_x += (len(ws) + 1) * char_w
|
|
else:
|
|
words.append({
|
|
"text": text,
|
|
"box": {"x": x_min, "y": y_min, "width": w, "height": h},
|
|
"confidence": score,
|
|
})
|
|
|
|
lines.append({
|
|
"text": text,
|
|
"box": {"x": x_min, "y": y_min, "width": w, "height": h},
|
|
"confidence": score,
|
|
"words": words,
|
|
"polygon": [{"x": pt[0], "y": pt[1]} for pt in box],
|
|
})
|
|
|
|
elapsed_ms = (time.time() - start_time) * 1000.0
|
|
return {
|
|
"imageWidth": width,
|
|
"imageHeight": height,
|
|
"lines": lines,
|
|
"processTimeMs": round(elapsed_ms, 2),
|
|
}
|
|
|
|
|
|
def _extract_page_font_spans(page: Any) -> List[Dict[str, Any]]:
|
|
"""
|
|
Extract text spans with font metadata from a PDF page using the C++ engine.
|
|
Returns a list of dicts: {x, y, width, height, fontName, fontSize}
|
|
Falls back to empty list if the engine doesn't support it.
|
|
"""
|
|
spans: List[Dict[str, Any]] = []
|
|
|
|
# Strategy 1: page.extract_document_model() — authoritative C++ engine model
|
|
# This is the same API used by the TextEditLayer for rich font information
|
|
if hasattr(page, "extract_document_model"):
|
|
try:
|
|
model = page.extract_document_model()
|
|
if model and hasattr(model, "paragraphs"):
|
|
for para in model.paragraphs:
|
|
for line in para.lines:
|
|
for run in line.runs:
|
|
fn = str(run.font_name) if hasattr(run, "font_name") and run.font_name else ""
|
|
fs = float(run.font_size) if hasattr(run, "font_size") and run.font_size else 0.0
|
|
x = float(run.x) if hasattr(run, "x") else 0.0
|
|
y = float(run.y) if hasattr(run, "y") else 0.0
|
|
w = float(run.w) if hasattr(run, "w") else 0.0
|
|
h = float(run.h) if hasattr(run, "h") and run.h > 0 else fs
|
|
is_bold = "bold" in fn.lower() or (hasattr(run, "flags") and bool(run.flags & 2))
|
|
if fn and fs > 0:
|
|
spans.append({
|
|
"x": x, "y": y, "width": w, "height": h,
|
|
"fontName": fn, "fontSize": fs, "isBold": is_bold,
|
|
})
|
|
except Exception as err:
|
|
print(f"[OCR] Strategy 1 error: {err}")
|
|
|
|
# Strategy 2: page.extract_text_model() — JSON structured model
|
|
if not spans and hasattr(page, "extract_text_model"):
|
|
try:
|
|
import json
|
|
raw = page.extract_text_model()
|
|
model = json.loads(raw) if isinstance(raw, str) else raw
|
|
for para in (model.get("paragraphs") or []):
|
|
for line in (para.get("lines") or []):
|
|
for run in (line.get("runs") or []):
|
|
fn = run.get("font_name") or run.get("fontName") or ""
|
|
fs = float(run.get("font_size") or run.get("fontSize") or 0)
|
|
x = float(run.get("x", 0))
|
|
y = float(run.get("y", 0))
|
|
w = float(run.get("w", 0))
|
|
h = float(run.get("h", fs))
|
|
if fn and fs > 0:
|
|
spans.append({"x": x, "y": y, "width": w, "height": h,
|
|
"fontName": fn, "fontSize": fs})
|
|
except Exception:
|
|
pass
|
|
|
|
# Strategy 3: page.extract_display_list() — lower-level ops
|
|
if not spans and hasattr(page, "extract_display_list"):
|
|
try:
|
|
import json
|
|
dl_str = page.extract_display_list()
|
|
dl_ops = json.loads(dl_str) if dl_str else []
|
|
current_font = "Helvetica"
|
|
current_size = 12.0
|
|
for op in dl_ops:
|
|
op_name = op.get("op", "")
|
|
args = op.get("args", [])
|
|
if op_name == "Tf" and len(args) >= 2:
|
|
current_font = str(args[0]).lstrip("/")
|
|
try:
|
|
current_size = float(args[1])
|
|
except (ValueError, TypeError):
|
|
pass
|
|
elif op_name in ("Tj", "TJ", "'", "\"") and current_font:
|
|
spans.append({
|
|
"x": 0, "y": 0, "width": 100, "height": current_size,
|
|
"fontName": current_font, "fontSize": current_size,
|
|
})
|
|
except Exception:
|
|
pass
|
|
|
|
return spans
|
|
|
|
|
|
def _match_font_to_box(box: Dict[str, float], spans: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Return the font span whose bounding box overlaps most with `box`.
|
|
Falls back to None if no spans available.
|
|
"""
|
|
if not spans:
|
|
return None
|
|
|
|
bx, by, bw, bh = box["x"], box["y"], box["width"], box["height"]
|
|
|
|
best = None
|
|
best_score = -1.0
|
|
|
|
for sp in spans:
|
|
# Intersection area
|
|
ix1 = max(bx, sp["x"])
|
|
iy1 = max(by, sp["y"])
|
|
ix2 = min(bx + bw, sp["x"] + sp["width"])
|
|
iy2 = min(by + bh, sp["y"] + sp["height"])
|
|
inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
|
|
if inter > best_score:
|
|
best_score = inter
|
|
best = sp
|
|
|
|
# If no overlap at all, fall back to closest span by vertical midpoint distance
|
|
if best_score <= 0:
|
|
mid_y = by + bh / 2.0
|
|
best = min(spans, key=lambda s: abs((s["y"] + s["height"] / 2.0) - mid_y))
|
|
|
|
return best
|
|
|
|
|
|
def process_pdf_page_ocr(doc: Any, page_index: int, dpi: int = 200) -> Dict[str, Any]:
|
|
if not is_ocr_available():
|
|
raise RuntimeError("RapidOCR engine is not available")
|
|
|
|
page = doc.get_page(page_index)
|
|
pdf_w, pdf_h = page.width, page.height
|
|
|
|
# ── Font span extraction ─────────────────────────────────────────────────
|
|
print(f"[OCR DEBUG] doc methods: {[m for m in dir(doc) if not m.startswith('_')]}")
|
|
print(f"[OCR DEBUG] page methods: {[m for m in dir(page) if not m.startswith('_')]}")
|
|
font_spans = _extract_page_font_spans(page)
|
|
print(f"[OCR] Page {page_index}: extracted {len(font_spans)} font spans from page")
|
|
if font_spans:
|
|
unique_fonts = set(s["fontName"] for s in font_spans)
|
|
print(f"[OCR] Unique fonts found: {unique_fonts}")
|
|
|
|
# ── Fallback: page-level and document-level fonts (for scanned PDFs) ──────
|
|
doc_font_name = None
|
|
if not font_spans:
|
|
# Try page.get_fonts() first
|
|
try:
|
|
if hasattr(page, "get_fonts"):
|
|
page_fonts = page.get_fonts()
|
|
print(f"[OCR DEBUG] page.get_fonts() returned {len(page_fonts) if page_fonts else 0} fonts")
|
|
if page_fonts:
|
|
import re as _re
|
|
for f in page_fonts:
|
|
fn = getattr(f, "font_name", "") or ""
|
|
nf = getattr(f, "normalized_family", "") or ""
|
|
clean = _re.sub(r"^[A-Z]{6}\+", "", fn).strip()
|
|
family = nf or clean
|
|
print(f"[OCR DEBUG] page_font item: fn='{fn}', nf='{nf}', clean='{clean}'")
|
|
if family and not any(x in family.lower() for x in ["symbol", "zapf", "wingding"]):
|
|
doc_font_name = family
|
|
print(f"[OCR] Page-level font fallback: {doc_font_name} (from {fn})")
|
|
break
|
|
except Exception as e:
|
|
print(f"[OCR] page.get_fonts() failed: {e}")
|
|
|
|
# Try doc.get_fonts(0, -1)
|
|
if not doc_font_name:
|
|
try:
|
|
if hasattr(doc, "get_fonts"):
|
|
doc_fonts = doc.get_fonts(0, -1)
|
|
print(f"[OCR DEBUG] doc.get_fonts(0, -1) returned {len(doc_fonts) if doc_fonts else 0} fonts")
|
|
if doc_fonts:
|
|
import re as _re
|
|
for f in doc_fonts:
|
|
fn = getattr(f, "font_name", "") or ""
|
|
nf = getattr(f, "normalized_family", "") or ""
|
|
clean = _re.sub(r"^[A-Z]{6}\+", "", fn).strip()
|
|
family = nf or clean
|
|
print(f"[OCR DEBUG] doc_font item: fn='{fn}', nf='{nf}', clean='{clean}'")
|
|
if family and not any(x in family.lower() for x in ["symbol", "zapf", "wingding"]):
|
|
doc_font_name = family
|
|
print(f"[OCR] Doc-level font fallback: {doc_font_name} (from {fn})")
|
|
break
|
|
except Exception as e:
|
|
print(f"[OCR] doc.get_fonts(0, -1) failed: {e}")
|
|
|
|
if not doc_font_name:
|
|
print("[OCR] No font detected at all, using Helvetica default")
|
|
|
|
# Render page PNG via pdfengine
|
|
img_result = page.render(dpi)
|
|
png_bytes = bytes(img_result.data)
|
|
|
|
img = Image.open(io.BytesIO(png_bytes))
|
|
img_w, img_h = img.size
|
|
|
|
ocr_result = recognize_image_bytes(png_bytes)
|
|
|
|
# Scale coordinates from image pixels to PDF points
|
|
scale_x = pdf_w / float(img_w)
|
|
scale_y = pdf_h / float(img_h)
|
|
|
|
scaled_lines = []
|
|
for line in ocr_result["lines"]:
|
|
l_box = line["box"]
|
|
pdf_l_box = {
|
|
"x": round(l_box["x"] * scale_x, 2),
|
|
"y": round(l_box["y"] * scale_y, 2),
|
|
"width": round(l_box["width"] * scale_x, 2),
|
|
"height": round(l_box["height"] * scale_y, 2),
|
|
}
|
|
|
|
# ── Font matching ────────────────────────────────────────────────────
|
|
matched = _match_font_to_box(pdf_l_box, font_spans)
|
|
if matched:
|
|
detected_font_name = matched["fontName"]
|
|
elif doc_font_name:
|
|
detected_font_name = doc_font_name
|
|
else:
|
|
detected_font_name = "Helvetica"
|
|
# Derive display name: strip subset prefix like "ABCDEF+"
|
|
import re as _re
|
|
detected_font_name = _re.sub(r"^[A-Z]{6}\+", "", detected_font_name).strip() or "Helvetica"
|
|
detected_font_size = float(matched["fontSize"]) if matched else round(pdf_l_box["height"] * 0.75, 2)
|
|
|
|
scaled_words = []
|
|
for w in line["words"]:
|
|
w_box = w["box"]
|
|
scaled_words.append({
|
|
"text": w["text"],
|
|
"box": {
|
|
"x": round(w_box["x"] * scale_x, 2),
|
|
"y": round(w_box["y"] * scale_y, 2),
|
|
"width": round(w_box["width"] * scale_x, 2),
|
|
"height": round(w_box["height"] * scale_y, 2),
|
|
},
|
|
"confidence": w["confidence"],
|
|
})
|
|
|
|
# Determine bold weight:
|
|
is_bold = False
|
|
if matched and matched.get("isBold"):
|
|
is_bold = True
|
|
elif pdf_l_box["height"] >= 16.5 or "bold" in detected_font_name.lower():
|
|
is_bold = True
|
|
elif line["text"].isupper() and len(line["text"].strip()) >= 3:
|
|
is_bold = True
|
|
|
|
scaled_lines.append({
|
|
"text": line["text"],
|
|
"box": pdf_l_box,
|
|
"confidence": line["confidence"],
|
|
"words": scaled_words,
|
|
"polygon": [
|
|
{"x": round(pt["x"] * scale_x, 2), "y": round(pt["y"] * scale_y, 2)}
|
|
for pt in line["polygon"]
|
|
],
|
|
"fontName": detected_font_name,
|
|
"fontSize": detected_font_size,
|
|
"isBold": is_bold,
|
|
})
|
|
|
|
if scaled_lines:
|
|
print(f"[OCR] First line font: '{scaled_lines[0].get('fontName')}' size={scaled_lines[0].get('fontSize')} isBold={scaled_lines[0].get('isBold')}")
|
|
|
|
return {
|
|
"pageIndex": page_index,
|
|
"pageWidth": pdf_w,
|
|
"pageHeight": pdf_h,
|
|
"lines": scaled_lines,
|
|
"processTimeMs": ocr_result["processTimeMs"],
|
|
}
|
|
|