Spaces:
Sleeping
Sleeping
File size: 2,285 Bytes
2a9e166 | 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 | """PDF/image loading and OCR.
Native (embedded) text is pulled straight out of the PDF when present;
pages without a usable text layer (scanned documents) are rendered to an
image and run through Tesseract instead.
"""
import os
import fitz # PyMuPDF
import pytesseract
from PIL import Image
MIN_EMBEDDED_TEXT_LEN = 20
RENDER_DPI = 300
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"}
def _ocr_image(img):
data = pytesseract.image_to_data(img, output_type=pytesseract.Output.DICT)
confidences = [int(c) for c in data.get("conf", []) if str(c).isdigit() and int(c) >= 0]
mean_conf = round(sum(confidences) / len(confidences), 1) if confidences else None
text = pytesseract.image_to_string(img)
return text, mean_conf
def _process_pdf(file_path):
pages = []
doc = fitz.open(file_path)
try:
for page_num, page in enumerate(doc, start=1):
embedded_text = page.get_text().strip()
if len(embedded_text) >= MIN_EMBEDDED_TEXT_LEN:
pages.append({
"page": page_num,
"text": embedded_text,
"source": "embedded",
"ocr_confidence": None,
})
continue
pix = page.get_pixmap(dpi=RENDER_DPI)
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
ocr_text, mean_conf = _ocr_image(img)
pages.append({
"page": page_num,
"text": ocr_text.strip(),
"source": "ocr",
"ocr_confidence": mean_conf,
})
finally:
doc.close()
return pages
def _process_image(file_path):
img = Image.open(file_path).convert("RGB")
text, mean_conf = _ocr_image(img)
return [{
"page": 1,
"text": text.strip(),
"source": "ocr",
"ocr_confidence": mean_conf,
}]
def extract_pages(file_path):
"""Return a list of {page, text, source, ocr_confidence} dicts for a file."""
ext = os.path.splitext(file_path)[1].lower()
if ext == ".pdf":
return _process_pdf(file_path)
if ext in IMAGE_EXTENSIONS:
return _process_image(file_path)
raise ValueError(f"Unsupported file type: {ext}")
|