pdf_processor.py

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.")