app.py

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)