Spaces:
Sleeping
Sleeping
File size: 4,425 Bytes
0e39d80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """Text extraction from image and PDF uploads (JPG, PNG, PDF only).
For images, uses EasyOCR via the ocr module.
For PDFs, renders up to PDF_MAX_PAGES at 150 DPI and runs OCR on page 1.
"""
from __future__ import annotations
from pathlib import Path
import cv2
import fitz
import numpy as np
from config import MAX_IMAGE_DIMENSION, PDF_MAX_PAGES
from ml_utils.ocr import OcrResult, get_full_text, ocr_multipass
MIN_NATIVE_CHARS = 50
_PDF_DPI = 150
def _cap_image_dimension(img: np.ndarray, max_dim: int = MAX_IMAGE_DIMENSION) -> np.ndarray:
h, w = img.shape[:2]
longest = max(h, w)
if longest <= max_dim:
return img
scale = max_dim / longest
new_w = max(1, int(w * scale))
new_h = max(1, int(h * scale))
return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
def extract_plain_text(file_path: str) -> tuple[str, str]:
"""Returns (text, source) where source is native|ocr|failed."""
path = Path(file_path)
ext = path.suffix.lower()
try:
if ext == ".pdf":
return _extract_pdf(path)
if ext in (".jpg", ".jpeg", ".png"):
img = cv2.imread(str(path))
if img is None:
return "", "failed"
img = _cap_image_dimension(img)
return _ocr_image(img)
except Exception:
return "", "failed"
return "", "failed"
def extract_with_ocr_results(file_path: str) -> tuple[str, str, list[OcrResult], np.ndarray | None]:
"""Extract text AND return OCR results + image for the full pipeline.
Returns: (text, text_source, ocr_results, image_bgr_or_None)
"""
path = Path(file_path)
ext = path.suffix.lower()
try:
if ext in (".jpg", ".jpeg", ".png"):
img = cv2.imread(str(path))
if img is None:
return "", "failed", [], None
img = _cap_image_dimension(img)
results = ocr_multipass(img)
text = get_full_text(results)
return text, "ocr", results, img
if ext == ".pdf":
native_text = _get_pdf_text(path)
img = _pdf_page_to_bgr(path, page=0)
if img is not None:
img = _cap_image_dimension(img)
if len(native_text.strip()) >= MIN_NATIVE_CHARS:
results = ocr_multipass(img) if img is not None else []
return native_text, "native", results, img
if img is not None:
results = ocr_multipass(img)
text = get_full_text(results)
return text, "ocr", results, img
return native_text, "native" if native_text.strip() else "failed", [], None
except Exception:
pass
return "", "failed", [], None
def _ocr_image(img: np.ndarray) -> tuple[str, str]:
"""Run EasyOCR multipass on an image."""
results = ocr_multipass(img)
text = get_full_text(results)
if text.strip():
return text, "ocr"
return "", "failed"
def _get_pdf_text(path: Path) -> str:
"""Get native text from first PDF_MAX_PAGES pages."""
doc = fitz.open(str(path))
text_parts: list[str] = []
for page_num in range(min(len(doc), PDF_MAX_PAGES)):
text_parts.append(doc.load_page(page_num).get_text())
doc.close()
return "\n".join(text_parts).strip()
def _extract_pdf(path: Path) -> tuple[str, str]:
text = _get_pdf_text(path)
if len(text) >= MIN_NATIVE_CHARS:
return text, "native"
img = _pdf_page_to_bgr(path, page=0)
if img is not None:
img = _cap_image_dimension(img)
return _ocr_image(img)
return text, "native" if text else "failed"
def pdf_to_image_bgr(file_path: str, page: int = 0) -> np.ndarray | None:
return _pdf_page_to_bgr(Path(file_path), page)
def _pdf_page_to_bgr(path: Path, page: int = 0) -> np.ndarray | None:
try:
doc = fitz.open(str(path))
if page >= len(doc):
doc.close()
return None
zoom = _PDF_DPI / 72.0
pix = doc.load_page(page).get_pixmap(matrix=fitz.Matrix(zoom, zoom))
doc.close()
img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
if pix.n == 4:
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
elif pix.n == 3:
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
return img
except Exception:
return None
|