Spaces:
Running
Running
File size: 4,086 Bytes
b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 bd7234d b611f38 | 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 | """
PDF and Image handling utilities for OCR space.
Supports multi-page PDF rendering from paths or in-memory bytes with guaranteed file handle closure.
"""
import os
from typing import List, Union, Tuple, Optional
from PIL import Image
import io
import logging
logger = logging.getLogger("PDFUtils")
def is_pdf_file(file_input: Union[str, bytes, None]) -> bool:
"""Check if the input represents a PDF."""
if isinstance(file_input, str):
return file_input.lower().endswith(".pdf")
if isinstance(file_input, bytes):
return file_input.startswith(b"%PDF")
return False
def convert_pdf_bytes_to_images(pdf_bytes: bytes, max_pages: int = 10, dpi: int = 200) -> List[Image.Image]:
"""Converts in-memory PDF bytes into a list of PIL Images (one per page)."""
images = []
# 1. Try pypdfium2 (in-memory, no temp file)
try:
import pypdfium2 as pdfium
pdf = pdfium.PdfDocument(pdf_bytes)
try:
total_pages = min(len(pdf), max_pages)
scale = dpi / 72.0
for i in range(total_pages):
page = pdf[i]
image = page.render(scale=scale).to_pil()
images.append(image.convert("RGB"))
if images:
return images
finally:
pdf.close()
except Exception as e:
logger.debug(f"pypdfium2 byte render fallback: {e}")
# 2. Try fitz (PyMuPDF in-memory)
try:
import fitz
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
try:
total_pages = min(len(doc), max_pages)
for i in range(total_pages):
page = doc[i]
pix = page.get_pixmap(dpi=dpi)
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
images.append(img)
if images:
return images
finally:
doc.close()
except Exception as e:
logger.debug(f"fitz byte render fallback: {e}")
if not images:
raise RuntimeError("Could not decode in-memory PDF with pypdfium2 or fitz.")
return images
def convert_pdf_to_images(pdf_path: str, max_pages: int = 10, dpi: int = 200) -> List[Image.Image]:
"""Converts a PDF file from a path into a list of PIL Images."""
if not os.path.exists(pdf_path):
raise FileNotFoundError(f"PDF not found: {pdf_path}")
with open(pdf_path, "rb") as f:
pdf_bytes = f.read()
return convert_pdf_bytes_to_images(pdf_bytes, max_pages=max_pages, dpi=dpi)
def load_input_image_or_pdf(
file_input: Union[str, bytes, Image.Image, None], page_index: int = 0
) -> Tuple[Optional[Image.Image], int, List[Image.Image]]:
"""
Helper to process upload file (image, PDF path, or PDF bytes) and return:
- (selected_image, total_pages, all_page_images)
"""
if file_input is None:
return None, 0, []
if isinstance(file_input, Image.Image):
return file_input.convert("RGB"), 1, [file_input.convert("RGB")]
if isinstance(file_input, bytes):
try:
img = Image.open(io.BytesIO(file_input)).convert("RGB")
return img, 1, [img]
except Exception:
pages = convert_pdf_bytes_to_images(file_input)
total = len(pages)
if total == 0:
raise ValueError("PDF contains no readable pages.")
idx = max(0, min(page_index, total - 1))
return pages[idx], total, pages
if isinstance(file_input, str):
if not os.path.exists(file_input):
raise FileNotFoundError(f"File not found: {file_input}")
if is_pdf_file(file_input):
pages = convert_pdf_to_images(file_input)
total = len(pages)
if total == 0:
raise ValueError("PDF contains no readable pages.")
idx = max(0, min(page_index, total - 1))
return pages[idx], total, pages
else:
img = Image.open(file_input).convert("RGB")
return img, 1, [img]
return None, 0, []
|