VOR / app.py
Preeeeet's picture
Update app.py
21e4481 verified
Raw
History Blame Contribute Delete
4.12 kB
import gradio as gr
import os
import tempfile
from PyPDF2 import PdfReader, PdfWriter
from reportlab.pdfgen import canvas
from PIL import Image, ImageDraw
import shutil
class MultiDocMerger:
def __init__(self):
self.temp_dir = tempfile.mkdtemp()
self.documents = []
self.pages = []
def add_document(self, file):
file_path = os.path.join(self.temp_dir, file.name)
with open(file_path, 'wb') as f:
f.write(file.file.read()) # Use file.file.read() instead of file.read()
self.documents.append(file_path)
self._extract_pages(file_path)
return self._generate_preview()
self.documents.append(file_path)
self._extract_pages(file_path)
return self._generate_preview()
def _extract_pages(self, file_path):
reader = PdfReader(file_path)
for page_num, page in enumerate(reader.pages):
temp_pdf = os.path.join(self.temp_dir, f"page_{len(self.pages)}.pdf")
writer = PdfWriter()
writer.add_page(page)
with open(temp_pdf, 'wb') as f:
writer.write(f)
self.pages.append(temp_pdf)
def rearrange_documents(self, order):
self.documents = [self.documents[i] for i in order]
return self._generate_preview()
def rearrange_pages(self, order):
self.pages = [self.pages[i] for i in order]
return self._generate_preview()
def delete_page(self, index):
if 0 <= index < len(self.pages):
del self.pages[index]
return self._generate_preview()
def add_blank_page(self):
blank_path = os.path.join(self.temp_dir, f"blank_{len(self.pages)}.pdf")
can = canvas.Canvas(blank_path)
can.showPage()
can.save()
self.pages.append(blank_path)
return self._generate_preview()
def generate_index(self):
index_content = "Index Page\n\n"
for i, doc in enumerate(self.documents):
index_content += f"{i + 1}. {os.path.basename(doc)}\n"
index_pdf = os.path.join(self.temp_dir, "index.pdf")
can = canvas.Canvas(index_pdf)
can.drawString(100, 750, index_content)
can.save()
self.pages.insert(0, index_pdf)
def export_merged_pdf(self):
self.generate_index()
output = PdfWriter()
for page in self.pages:
reader = PdfReader(page)
output.add_page(reader.pages[0])
output_path = os.path.join(self.temp_dir, 'merged_output.pdf')
with open(output_path, 'wb') as f:
output.write(f)
return output_path
def _generate_preview(self):
return [os.path.basename(page) for page in self.pages]
def cleanup(self):
shutil.rmtree(self.temp_dir)
merger = MultiDocMerger()
with gr.Blocks() as demo:
gr.Markdown("# Multi-Document Merger and Organizer")
with gr.Column():
file_upload = gr.File(label="Upload Files", file_count="multiple")
preview_gallery = gr.Gallery(label="Document Preview and Reordering", interactive=True)
delete_index = gr.Number(label="Delete Page Index")
blank_button = gr.Button("Add Blank Page")
merge_button = gr.Button("Merge and Export PDF")
output_file = gr.File()
def upload_files(files):
paths = []
for file in files:
paths = merger.add_document(file)
return paths
def merge_and_export():
return merger.export_merged_pdf()
def delete_page(index):
return merger.delete_page(index)
def add_blank_page():
return merger.add_blank_page()
def reorder_pages(order):
return merger.rearrange_pages(order)
file_upload.upload(upload_files, inputs=file_upload, outputs=preview_gallery)
delete_index.submit(delete_page, inputs=delete_index, outputs=preview_gallery)
blank_button.click(add_blank_page, outputs=preview_gallery)
preview_gallery.change(reorder_pages, inputs=preview_gallery, outputs=preview_gallery)
merge_button.click(merge_and_export, outputs=output_file)
demo.launch()