app.py

import os
import gradio as gr
import tempfile
import shutil
import pandas as pd
from PIL import Image
from preprocess import convert_pdf_to_images, preprocess_image
from llm_utils import load_image, generate_response

# Global temporary directory for image processing
temp_dir = tempfile.mkdtemp()

pdf_image_map = {}  # Map from PDF to its image list
image_preview_map = {}  # Map from image name to full path


def extract_images_from_pdfs(pdf_files):
    global pdf_image_map, image_preview_map
    pdf_image_map.clear()
    image_preview_map.clear()

    previews = []

    for pdf in pdf_files:
        image_paths = convert_pdf_to_images(pdf.name, temp_dir)
        pdf_image_map[pdf.name] = image_paths

        for img_path in image_paths:
            img_name = os.path.basename(img_path)
            image_preview_map[img_name] = img_path
            previews.append((img_name, img_path))

    # Return preview (tuples: filename, image_path)
    return [img for _, img in previews]


def process_selected_images(selected_images, output_excel_name):
    results = []

    for img_name in selected_images:
        img_path = image_preview_map.get(img_name)
        if img_path is None:
            continue

        processed = preprocess_image(img_path)
        processed_path = os.path.join(temp_dir, f"processed_{img_name}")
        Image.fromarray(processed).save(processed_path)

        # Run LLM
        pixel_values = load_image(processed_path)
        response = generate_response(pixel_values)

        # Clean filename to find source pdf
        base_name = img_name.rsplit("_page_", 1)[0] + ".pdf"
        results.append({"Source PDF": base_name, "Page Image": img_name, "LLM Output": response})

    df = pd.DataFrame(results)
    output_excel = os.path.join(temp_dir, output_excel_name)
    df.to_excel(output_excel, index=False)
    return output_excel


def reset_all():
    shutil.rmtree(temp_dir)
    os.makedirs(temp_dir, exist_ok=True)


with gr.Blocks() as demo:
    gr.Markdown("## 🧠 PDF → Image → LLM → Excel Output")

    with gr.Row():
        pdf_input = gr.File(file_types=[".pdf"], file_count="multiple", label="📎 Upload multiple PDF files")
        extract_btn = gr.Button("🔍 Extract Images")

    gallery = gr.Gallery(label="🖼️ Choose images to process", columns=4, allow_preview=True, interactive=True, show_label=True).style(grid=4)
    selected_images = gr.CheckboxGroup(choices=[], label="Select image filenames for processing")

    with gr.Row():
        output_name = gr.Textbox(label="📄 Output Excel filename", value="output.xlsx")
        generate_btn = gr.Button("🚀 Generate Excel")

    excel_output = gr.File(label="📥 Download Excel")

    def update_gallery(pdf_files):
        previews = extract_images_from_pdfs(pdf_files)
        choices = list(image_preview_map.keys())
        return gr.update(value=previews), gr.update(choices=choices, value=[])

    extract_btn.click(update_gallery, inputs=[pdf_input], outputs=[gallery, selected_images])
    generate_btn.click(fn=process_selected_images, inputs=[selected_images, output_name], outputs=[excel_output])

if __name__ == "__main__":
    demo.launch()