Create app.py

app.py

@@ -0,0 +1,161 @@
+import os
+from PyPDF2 import PdfReader, PdfWriter
+from io import BytesIO
+import gradio as gr
+
+def get_pdf_size(pdf_writer):
+    """获取当前PdfWriter对象的内容大小（以字节为单位）"""
+    buffer = BytesIO()
+    pdf_writer.write(buffer)
+    return buffer.tell()
+
+def split_pdf_by_pages(input_pdf_path, output_folder, split_pages):
+    """按分页数分割PDF"""
+    # 创建输出文件夹
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    # 读取PDF文件
+    reader = PdfReader(input_pdf_path)
+    total_pages = len(reader.pages)
+
+    # 检查用户输入的分页数是否有效
+    if any(page > total_pages or page < 1 for page in split_pages):
+        return f"错误：分页数超出PDF总页数范围（总页数：{total_pages}）。"
+
+    # 对分页数进行排序并去重
+    split_pages = sorted(set(split_pages))
+
+    # 切割PDF
+    start_page = 0
+    result_files = []
+    for i, split_page in enumerate(split_pages):
+        writer = PdfWriter()
+        end_page = split_page - 1  # PyPDF2的页码从0开始
+
+        # 添加从 start_page 到 end_page 的页面
+        for page_num in range(start_page, end_page + 1):
+            writer.add_page(reader.pages[page_num])
+
+        # 保存切割后的PDF文件
+        output_pdf_path = os.path.join(output_folder, f"part_{i+1}.pdf")
+        with open(output_pdf_path, "wb") as output_pdf:
+            writer.write(output_pdf)
+
+        result_files.append(output_pdf_path)
+        start_page = end_page + 1
+
+    # 处理最后一组页面
+    if start_page < total_pages:
+        writer = PdfWriter()
+        for page_num in range(start_page, total_pages):
+            writer.add_page(reader.pages[page_num])
+
+        output_pdf_path = os.path.join(output_folder, f"part_{len(split_pages) + 1}.pdf")
+        with open(output_pdf_path, "wb") as output_pdf:
+            writer.write(output_pdf)
+
+        result_files.append(output_pdf_path)
+
+    return result_files
+
+def split_pdf_by_size(input_pdf_path, output_folder, max_size_mb):
+    """按文件大小分割PDF"""
+    # 创建输出文件夹
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    # 读取PDF文件
+    reader = PdfReader(input_pdf_path)
+    total_pages = len(reader.pages)
+
+    # 初始化变量
+    writer = PdfWriter()
+    current_part = 1
+    max_size_bytes = max_size_mb * 1024 * 1024  # 将MB转换为字节
+    result_files = []
+
+    # 逐页处理
+    for page_num in range(total_pages):
+        # 添加当前页
+        writer.add_page(reader.pages[page_num])
+
+        # 检查当前文件大小
+        current_size = get_pdf_size(writer)
+        if current_size >= max_size_bytes:
+            # 如果超过最大大小，保存当前部分
+            output_pdf_path = os.path.join(output_folder, f"part_{current_part}.pdf")
+            with open(output_pdf_path, "wb") as output_pdf:
+                writer.write(output_pdf)
+            result_files.append(output_pdf_path)
+
+            # 重置writer并增加部分计数
+            writer = PdfWriter()
+            current_part += 1
+
+    # 保存最后一部分（如果有剩余页面）
+    if len(writer.pages) > 0:
+        output_pdf_path = os.path.join(output_folder, f"part_{current_part}.pdf")
+        with open(output_pdf_path, "wb") as output_pdf:
+            writer.write(output_pdf)
+        result_files.append(output_pdf_path)
+
+    return result_files
+
+def process_pdf(input_pdf, mode, split_pages=None, max_size_mb=None):
+    """处理PDF文件"""
+    output_folder = "output_parts"
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    if mode == "按分页数分割":
+        if not split_pages:
+            return "错误：请输入分页数。"
+        split_pages = [int(page) for page in split_pages.split(",")]
+        result_files = split_pdf_by_pages(input_pdf, output_folder, split_pages)
+    elif mode == "按文件大小分割":
+        if not max_size_mb:
+            return "错误：请输入最大文件大小。"
+        result_files = split_pdf_by_size(input_pdf, output_folder, float(max_size_mb))
+    else:
+        return "错误：无效的模式。"
+
+    if isinstance(result_files, str):  # 如果返回的是错误信息
+        return result_files
+
+    # 返回所有生成的文件
+    return result_files
+
+# Gradio 界面
+with gr.Blocks() as demo:
+    gr.Markdown("# PDF 分割工具")
+    with gr.Row():
+        input_pdf = gr.File(label="上传PDF文件", type="filepath")
+        mode = gr.Radio(choices=["按分页数分割", "按文件大小分割"], label="选择分割模式")
+    with gr.Row():
+        split_pages = gr.Textbox(label="分页数（例如：3,5,10）", visible=True)
+        max_size_mb = gr.Number(label="每部分的最大大小（MB）", visible=False)
+    with gr.Row():
+        output_files = gr.Files(label="分割后的文件")
+    with gr.Row():
+        submit_btn = gr.Button("开始分割")
+
+    # 动态显示/隐藏输入框
+    def toggle_inputs(mode):
+        if mode == "按分页数分割":
+            return gr.Textbox(visible=True), gr.Number(visible=False)
+        else:
+            return gr.Textbox(visible=False), gr.Number(visible=True)
+
+    mode.change(toggle_inputs, inputs=mode, outputs=[split_pages, max_size_mb])
+
+    # 绑定处理函数
+    submit_btn.click(
+        process_pdf,
+        inputs=[input_pdf, mode, split_pages, max_size_mb],
+        outputs=output_files
+    )
+    theme=gr.themes.Soft()
+
+# 启动应用
+demo.launch()