Spaces:
Runtime error
Runtime error
File size: 3,175 Bytes
1b1359a a6c21d1 1b1359a 43facdb 1b1359a 43facdb 1b1359a | 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 | import os
import subprocess
# Temporarily force reinstall transformers for debugging
subprocess.run(["pip", "install", "--force-reinstall", "transformers==4.44.2"], check=True)
import gradio as gr
import shutil
from PIL import Image
from utils.pdf_to_images import convert_pdf_to_images, convert_image_to_pil
from utils.hunyuan_ocr import HunyuanOCR
from utils.pdf_builder import build_searchable_pdf
# Get Hugging Face token from environment variables (for gated models)
hf_token = os.environ.get("HF_TOKEN")
# Initialize OCR model
ocr_model = HunyuanOCR(token=hf_token)
def process_document(file_obj):
if file_obj is None:
return [], None, "Please upload a PDF or image file."
input_path = file_obj.name
file_extension = os.path.splitext(input_path)[1].lower()
images = []
if file_extension == ".pdf":
try:
images = convert_pdf_to_images(input_path)
except Exception as e:
return [], None, f"Error converting PDF to images: {e}"
elif file_extension in [".png", ".jpg", ".jpeg", ".tiff", ".tif"]:
try:
images = [convert_image_to_pil(input_path)]
except Exception as e:
return [], None, f"Error loading image: {e}"
else:
return [], None, "Unsupported file type. Please upload a PDF or image."
ocr_results = []
preview_images = []
for i, image in enumerate(images):
if max(image.width, image.height) > 10000:
return [], None, f"Page {i+1} is too large ({image.width}x{image.height}px). Max 10000px allowed."
try:
ocr_output = ocr_model.extract_ocr(image)
ocr_results.append(ocr_output)
preview_images.append(ocr_model.visualize_ocr(image.copy(), ocr_output))
except Exception as e:
return [], None, f"Error during OCR processing for page {i+1}: {e}"
output_pdf_path = "searchable_output.pdf"
try:
build_searchable_pdf(images, ocr_results, output_pdf_path)
except Exception as e:
return [], None, f"Error building searchable PDF: {e}"
return preview_images, output_pdf_path, "Searchable PDF generated successfully!"
# Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("# Searchable PDF Generator with Hunyuan-DiT")
gr.Markdown("Upload a PDF or image file to convert it into a layout-preserving, searchable PDF.")
with gr.Row():
file_upload = gr.File(label="Upload PDF or Image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".tif"])
convert_button = gr.Button("Convert to Searchable PDF")
status_message = gr.Textbox(label="Status", interactive=False)
ocr_preview_gallery = gr.Gallery(label="OCR Preview (Bounding Boxes)", columns=2, object_fit="contain", height="auto")
output_pdf_file = gr.File(label="Download Searchable PDF", interactive=False)
convert_button.click(
process_document,
inputs=[file_upload],
outputs=[ocr_preview_gallery, output_pdf_file, status_message]
)
if __name__ == "__main__":
demo.launch()
|