Spaces:
Sleeping
Sleeping
File size: 10,315 Bytes
bd1d88e 9c6f858 bd1d88e 8bc0dc9 bd1d88e 8bc0dc9 bd1d88e f3b29b4 bd1d88e 8bc0dc9 bd1d88e 8bc0dc9 bd1d88e 905a756 bd1d88e 905a756 bd1d88e 905a756 bd1d88e 905a756 bd1d88e 905a756 bd1d88e 905a756 bd1d88e 8bc0dc9 bd1d88e 8bc0dc9 bd1d88e | 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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 | 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.")
|