import os from PIL import Image from pdf2image import convert_from_path import pytesseract import fitz # PyMuPDF from diffusers import AutoPipelineForImage2Image # Corrected import import torch import gc # Import garbage collector import shutil # Initialize Stable Diffusion pipeline outside the function to avoid reloading for each call pipeline = None def _initialize_stable_diffusion_pipeline(): global pipeline if pipeline is None or pipeline == "placeholder_active": try: # Use a publicly available image-to-image model as an example. # For Stable Diffusion 3.5 Large, you would replace this model identifier # and potentially handle authentication (e.g., with a Hugging Face token). model_id = "stabilityai/sd-image-variations-diffusers" # Corrected model ID pipeline = AutoPipelineForImage2Image.from_pretrained(model_id, torch_dtype=torch.float16) if torch.cuda.is_available(): pipeline = pipeline.to("cuda") # Move pipeline to GPU if available print(f"Successfully loaded Stable Diffusion pipeline: {model_id}") except Exception as e: print(f"Could not load Stable Diffusion pipeline. Falling back to placeholder. Error: {e}") pipeline = "placeholder_active" # Indicate that placeholder should be used return pipeline def cleanup_sd_pipeline(): global pipeline if pipeline is not None and pipeline != "placeholder_active": try: del pipeline pipeline = None if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() print("Stable Diffusion pipeline resources explicitly cleaned up.") except Exception as e: print(f"Error during Stable Diffusion pipeline cleanup: {e}") def pdf_to_images(pdf_path, output_folder): """ Converts each page of a PDF into an image. Args: pdf_path (str): Path to the input PDF file. output_folder (str): Directory to save the output images. Returns: list: A list of paths to the generated image files. """ if not os.path.exists(output_folder): os.makedirs(output_folder) images = convert_from_path(pdf_path, dpi=300) # dpi for high resolution image_paths = [] for i, image in enumerate(images): image_path = os.path.join(output_folder, f"page_{i+1}.png") image.save(image_path, "PNG") image.close() # Explicitly close the PIL Image object after saving image_paths.append(image_path) return image_paths def enhance_image_with_stable_diffusion(image_path, output_folder): """ Enhances an image using a Stable Diffusion image-to-image model. If the model cannot be loaded, it falls back to simply copying the image. Args: image_path (str): Path to the input image file. output_folder (str): Directory to save the enhanced image. Returns: str: Path to the enhanced image file. """ _initialize_stable_diffusion_pipeline() if pipeline == "placeholder_active" or pipeline is None: print("Using placeholder for image enhancement (Stable Diffusion pipeline not loaded).") # Fallback to copying the image if the pipeline is not loaded enhanced_image_path = os.path.join(output_folder, os.path.basename(image_path)) with Image.open(image_path) as img: img.save(enhanced_image_path) return enhanced_image_path with Image.open(image_path).convert("RGB") as img: # The strength parameter controls how much the image is allowed to change. # A lower value means less change, a higher value means more creative freedom for the model. # You might need to experiment with this value. enhanced_img = pipeline(image=img, strength=0.75).images[0] enhanced_image_path = os.path.join(output_folder, "enhanced_" + os.path.basename(image_path)) enhanced_img.save(enhanced_image_path, "PNG") enhanced_img.close() # Explicitly close the output image too print(f"Image enhanced with Stable Diffusion: {enhanced_image_path}") return enhanced_image_path def perform_ocr(image_path, tesseract_config=""): """ Performs OCR on an image and returns the extracted text and its bounding box information. Args: image_path (str): Path to the input image file. tesseract_config (str): Optional Tesseract configuration string (e.g., '--psm 3 --oem 1'). Returns: dict: A dictionary containing OCR results, including text and bounding box data. """ with Image.open(image_path) as img: # Get OCR results, including bounding box information # output_type=pytesseract.Output.DICT gives a dictionary with box data ocr_data = pytesseract.image_to_data(img, output_type=pytesseract.Output.DICT, config=tesseract_config) return ocr_data def create_searchable_pdf(original_pdf_path, ocr_results_per_page, output_pdf_path, enhanced_image_paths): """ Creates a searchable PDF by overlaying OCR'd text invisibly on top of the original/enhanced image pages. Args: original_pdf_path (str): Path to the original input PDF file. ocr_results_per_page (list): A list of OCR results (dictionaries) for each page. output_pdf_path (str): Path to save the searchable PDF. enhanced_image_paths (list): List of paths to the enhanced images for each page. Returns: str: Path to the created searchable PDF. """ doc = fitz.open() original_doc = fitz.open(original_pdf_path) for i, ocr_data in enumerate(ocr_results_per_page): # Use the enhanced image for the background img_path = enhanced_image_paths[i] with Image.open(img_path) as img: img_width, img_height = img.size page = doc.new_page(width=img_width, height=img_height) # Add the enhanced image as the background page.insert_image(page.rect, filename=img_path) # Add invisible text layer for j in range(len(ocr_data['text'])): text = ocr_data['text'][j] left = ocr_data['left'][j] top = ocr_data['top'][j] width = ocr_data['width'][j] height = ocr_data['height'][j] # Ensure text is not empty and coordinates are valid if text.strip() and width > 0 and height > 0: # PyMuPDF coordinates are (x0, y0, x1, y1) # Tesseract's coordinates are (left, top, width, height) # We need to convert Tesseract's origin (top-left of image) to PyMuPDF's (top-left of page) # For a simple overlay, we can use the exact coordinates rect = fitz.Rect(left, top, left + width, top + height) # Insert text invisibly # Use a small font size and render mode 3 (invisible) so it doesn't appear visually # We approximate font size based on bounding box height for better searchability font_size = max(1, height * 0.8) # Heuristic to get a reasonable font size text_instance = page.insert_textbox( rect, text, fontname="helv", # A common font fontsize=font_size, render_mode=3, # Invisible text fill=(0, 0, 0) # Black color (doesn't matter much for invisible text) ) doc.save(output_pdf_path) doc.close() original_doc.close() return output_pdf_path if __name__ == '__main__': # Example usage (replace with your PDF path and desired output folder) # Make sure you have a sample.pdf in the same directory or provide its full path sample_pdf = "sample.pdf" output_image_folder = "output_images" enhanced_image_folder = "enhanced_images_sd" searchable_pdf_output = "searchable_sample_sd.pdf" # Create dummy sample.pdf if it doesn't exist for testing if not os.path.exists(sample_pdf): print(f"'{sample_pdf}' not found. Creating a dummy PDF for testing.") doc = fitz.open() page = doc.new_page() page.insert_textbox(fitz.Rect(50, 50, 200, 100), "Hello World!\nThis is a test PDF.") doc.save(sample_pdf) doc.close() print(f"Starting PDF processing for '{sample_pdf}'...") # Ensure directories are clean for a fresh run if os.path.exists(output_image_folder): shutil.rmtree(output_image_folder) if os.path.exists(enhanced_image_folder): shutil.rmtree(enhanced_image_folder) # Step 1: Convert PDF to images images = pdf_to_images(sample_pdf, output_image_folder) print(f"Generated {len(images)} images.") # Step 2: Enhance images with Stable Diffusion _initialize_stable_diffusion_pipeline() enhanced_images = [] for img_path in images: enhanced_path = enhance_image_with_stable_diffusion(img_path, enhanced_image_folder) enhanced_images.append(enhanced_path) print(f"Enhanced {len(enhanced_images)} images.") # Step 3: Perform OCR on enhanced images all_ocr_results = [] tess_config = r'--psm 3 --oem 1' for img_path in enhanced_images: ocr_data = perform_ocr(img_path, tesseract_config=tess_config) all_ocr_results.append(ocr_data) print(f"Performed OCR on {len(all_ocr_results)} enhanced images.") # Step 4: Create Searchable PDF create_searchable_pdf(sample_pdf, all_ocr_results, searchable_pdf_output, enhanced_images) print(f"Searchable PDF created at: {searchable_pdf_output}") # Clean up created folders after example run if os.path.exists(output_image_folder): shutil.rmtree(output_image_folder) if os.path.exists(enhanced_image_folder): shutil.rmtree(enhanced_image_folder) # Clean up Stable Diffusion pipeline after example run cleanup_sd_pipeline() print("Processing complete.")