Spaces:
Running
Running
| import os | |
| import tempfile | |
| import logging | |
| from typing import List | |
| import math | |
| import gradio as gr | |
| import requests | |
| from PIL import Image | |
| from pdf2image import convert_from_path, convert_from_bytes | |
| from pdf2image.exceptions import PDFInfoNotInstalledError, PDFPageCountError | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") | |
| logger = logging.getLogger(__name__) | |
| def stitch_images_vertically(images: List[Image.Image]) -> Image.Image: | |
| if not images: | |
| return None | |
| if not all(isinstance(i, Image.Image) for i in images): | |
| logger.error("Non-Image object found in list for vertical stitching.") | |
| return None | |
| max_width = max(img.width for img in images) | |
| total_height = sum(img.height for img in images) | |
| stitched_image = Image.new('RGB', (max_width, total_height), (255, 255, 255)) | |
| current_y = 0 | |
| for img in images: | |
| stitched_image.paste(img, (0, current_y)) | |
| current_y += img.height | |
| return stitched_image | |
| def stitch_images_in_grid(images: List[Image.Image], num_columns: int, page_order: str) -> Image.Image: | |
| if not images: | |
| return None | |
| if page_order == "Top-to-Bottom (down)": | |
| num_images = len(images) | |
| num_rows = math.ceil(num_images / num_columns) | |
| columns = [images[i*num_rows : (i+1)*num_rows] for i in range(num_columns)] | |
| else: # Default to "Left-to-Right (across)" | |
| columns = [images[i::num_columns] for i in range(num_columns)] | |
| stitched_columns = [stitch_images_vertically(col) for col in columns if col] | |
| if not stitched_columns: | |
| return None | |
| max_height = max(col.height for col in stitched_columns if col) | |
| total_width = sum(col.width for col in stitched_columns if col) | |
| grid_image = Image.new('RGB', (total_width, max_height), (255, 255, 255)) | |
| current_x = 0 | |
| for col_img in stitched_columns: | |
| if col_img: | |
| grid_image.paste(col_img, (current_x, 0)) | |
| current_x += col_img.width | |
| return grid_image | |
| def process_pdf(pdf_file, pdf_url, dpi, num_columns, num_images, crop_top, crop_bottom, crop_left, crop_right, hide_annotations, page_order, progress=gr.Progress()): | |
| pdf_input_source = None | |
| is_bytes = False | |
| source_name = "document" | |
| progress(0, desc="Validating input...") | |
| if pdf_file is not None: | |
| logger.info(f"Processing uploaded file: {pdf_file.name}") | |
| pdf_input_source = pdf_file.name | |
| source_name = os.path.splitext(os.path.basename(pdf_file.name))[0] | |
| elif pdf_url and pdf_url.strip(): | |
| url = pdf_url.strip() | |
| logger.info(f"Processing file from URL: {url}") | |
| progress(0.1, desc="Downloading PDF from URL...") | |
| try: | |
| response = requests.get(url, timeout=45) | |
| response.raise_for_status() | |
| pdf_input_source = response.content | |
| source_name = os.path.splitext(os.path.basename(url.split('?')[0]))[0] | |
| is_bytes = True | |
| except requests.RequestException as e: | |
| raise gr.Error(f"Failed to download PDF from URL. Error: {e}") | |
| else: | |
| raise gr.Error("Please upload a PDF file or provide a valid URL.") | |
| progress(0.3, desc="Converting PDF pages to images...") | |
| logger.info(f"Using DPI: {dpi}, Hide Annotations: {hide_annotations}") | |
| try: | |
| if is_bytes: | |
| images = convert_from_bytes(pdf_input_source, dpi=dpi, hide_annotations=hide_annotations) | |
| else: | |
| images = convert_from_path(pdf_input_source, dpi=dpi, hide_annotations=hide_annotations) | |
| except (PDFInfoNotInstalledError, FileNotFoundError): | |
| raise gr.Error("Server configuration error: Poppler dependency is missing.") | |
| except (PDFPageCountError, Exception) as e: | |
| raise gr.Error(f"Failed to process the PDF. It might be corrupted or password-protected. Error: {e}") | |
| if not images: | |
| raise gr.Error("Could not extract any pages from the PDF. The file might be empty or invalid.") | |
| logger.info(f"Successfully converted {len(images)} pages to images.") | |
| cropped_images = [] | |
| if crop_top > 0 or crop_bottom > 0 or crop_left > 0 or crop_right > 0: | |
| progress(0.6, desc="Cropping images...") | |
| for i, img in enumerate(images): | |
| width, height = img.size | |
| left, top, right, bottom = crop_left, crop_top, width - crop_right, height - crop_bottom | |
| if left >= right or top >= bottom: | |
| raise gr.Error(f"Crop values are too large for page {i+1}. The page dimensions are {width}x{height}, but crop settings result in an invalid area.") | |
| cropped_images.append(img.crop((left, top, right, bottom))) | |
| else: | |
| cropped_images = images | |
| progress(0.7, desc=f"Splitting {len(cropped_images)} pages into {num_images} image(s)...") | |
| total_pages = len(cropped_images) | |
| effective_num_images = min(num_images, total_pages) | |
| chunk_size = math.ceil(total_pages / effective_num_images) | |
| image_chunks = [cropped_images[i:i + chunk_size] for i in range(0, total_pages, chunk_size)] | |
| output_paths = [] | |
| num_chunks = len(image_chunks) | |
| for i, chunk in enumerate(image_chunks): | |
| progress(0.75 + (0.2 * (i / num_chunks)), desc=f"Stitching image {i+1} of {num_chunks}...") | |
| if not chunk: | |
| continue | |
| if num_columns > 1: | |
| stitched_image = stitch_images_in_grid(chunk, num_columns, page_order) | |
| else: | |
| stitched_image = stitch_images_vertically(chunk) | |
| if stitched_image is None: | |
| logger.warning(f"Image stitching failed for chunk {i+1}.") | |
| continue | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{source_name}_stitched_{i+1}_") as tmp_file: | |
| stitched_image.save(tmp_file.name, "PNG") | |
| output_paths.append(tmp_file.name) | |
| if not output_paths: | |
| raise gr.Error("Image stitching failed for all pages.") | |
| logger.info(f"Final images saved to temporary paths: {output_paths}") | |
| progress(1, desc="Done!") | |
| return output_paths, output_paths | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown( | |
| """ | |
| # PDF Page Stitcher 📄 ➡️ 🖼️ | |
| Upload a PDF file or provide a URL. This tool will convert every page of the PDF into an image | |
| and then append them to create a single image that you can download. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| with gr.Tabs(): | |
| with gr.TabItem("Upload PDF"): | |
| pdf_file_input = gr.File(label="Upload PDF File", file_types=[".pdf"]) | |
| with gr.TabItem("From URL"): | |
| pdf_url_input = gr.Textbox(label="PDF URL", placeholder="e.g., https://arxiv.org/pdf/1706.03762.pdf") | |
| dpi_slider = gr.Slider(minimum=100, maximum=600, step=5, value=200, label="Image Resolution (DPI)") | |
| columns_slider = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Columns") | |
| num_images_slider = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Output Images", info="Splits the PDF pages into multiple output images.") | |
| with gr.Accordion("Advanced Options", open=False): | |
| hide_annotations_toggle = gr.Checkbox(value=True, label="Hide PDF Annotations (Links/Highlights)", info="Turn this on to remove the colored boxes that can appear around links and references.") | |
| page_order_radio = gr.Radio(["Left-to-Right (across)", "Top-to-Bottom (down)"], value="Left-to-Right (across)", label="Multi-Column Page Order", info="Determines how pages fill the columns.") | |
| with gr.Row(): | |
| crop_left = gr.Slider(minimum=0, maximum=500, step=10, value=0, label="Crop Left") | |
| crop_right = gr.Slider(minimum=0, maximum=500, step=10, value=0, label="Crop Right") | |
| with gr.Row(): | |
| crop_top = gr.Slider(minimum=0, maximum=500, step=10, value=0, label="Crop Top") | |
| crop_bottom = gr.Slider(minimum=0, maximum=500, step=10, value=0, label="Crop Bottom") | |
| submit_btn = gr.Button("Stitch PDF Pages", variant="primary") | |
| with gr.Column(scale=2): | |
| gr.Markdown("## Output") | |
| output_image_preview = gr.Gallery(label="Stitched Image(s) Preview", interactive=False, height=600) | |
| output_image_download = gr.File(label="Download Stitched Image(s)", interactive=False, file_count="multiple") | |
| submit_btn.click( | |
| fn=process_pdf, | |
| inputs=[ | |
| pdf_file_input, | |
| pdf_url_input, | |
| dpi_slider, | |
| columns_slider, | |
| num_images_slider, | |
| crop_top, | |
| crop_bottom, | |
| crop_left, | |
| crop_right, | |
| hide_annotations_toggle, | |
| page_order_radio | |
| ], | |
| outputs=[output_image_preview, output_image_download] | |
| ) | |
| demo.launch(server_name="0.0.0.0", server_port=7860, debug=True) |