Spaces:
Running
Running
| """ | |
| 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, [] | |