Update app.py

app.py

@@ -1,150 +1,127 @@
-import fitz  # PyMuPDF
-import cv2
-import numpy as np
 import os
+import cv2
 import gradio as gr
-import gc
+import numpy as np
 import zipfile
-from PIL import Image
-import tempfile
+import fitz  # PyMuPDF
+import gc
 
-# 常量設定
-MAX_PAGES = 250
-BG_PATH = "background.jpg"
+# 參數設定
 FG_DIR = "foregrounds"
 OUTPUT_DIR = "output"
-ZIP_PATH = "output/results.zip"
+ZIP_PATH = os.path.join(OUTPUT_DIR, "results.zip")
+BG_PATH = "background.jpg"
 TARGET_WIDTH, TARGET_HEIGHT = 2133, 1200
 POS1 = (1032, 0)
 POS2 = (4666, 0)
 
-# 初始化資料夾與權限檢查
-def check_write_permission(path):
+# 建立資料夾
+def ensure_dir(path):
     if not os.path.exists(path):
-        try:
-            os.makedirs(path)
-        except Exception as e:
-            raise RuntimeError(f"❌ 無法建立目錄 '{path}'：{e}")
-    elif not os.access(path, os.W_OK):
-        raise RuntimeError(f"❌ 目錄 '{path}' 存在，但沒有寫入權限！")
+        os.makedirs(path)
+    if not os.access(path, os.W_OK):
+        raise RuntimeError(f"❌ 沒有寫入權限：{path}")
 
-check_write_permission(FG_DIR)
-check_write_permission(OUTPUT_DIR)
+ensure_dir(FG_DIR)
+ensure_dir(OUTPUT_DIR)
 
-# PDF轉圖片（使用 PyMuPDF）
+# 將 PDF 每頁轉為圖檔
 def pdf_to_images(pdf_file):
-    try:
-        doc = fitz.open(pdf_file.name)  # 直接使用暫存檔案路徑
-        image_paths = []
-
-        for i, page in enumerate(doc, start=1):
-            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
-            img_path = os.path.join(FG_DIR, f"page_{i:03d}.jpg")
-            pix.save(img_path)
-            image_paths.append(img_path)
-
-        return image_paths
-    except Exception as e:
-        raise gr.Error(f"PDF轉換失敗: {str(e)}")
-
-# 合成圖片
+    doc = fitz.open(pdf_file.name)
+    paths = []
+    for i, page in enumerate(doc, start=1):
+        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
+        img_path = os.path.join(FG_DIR, f"page_{i:03d}.jpg")
+        pix.save(img_path)
+        paths.append(img_path)
+    return paths
+
+# 圖像合成功能
 def overlay_image(bg, fg, x, y):
     h, w = fg.shape[:2]
-    bh, bw = bg.shape[:2]
-    if x + w > bw or y + h > bh:
-        raise ValueError(f"前景圖超出背景邊界：({x},{y},{w},{h}) > 背景({bw},{bh})")
-
     if fg.shape[2] == 4:
         alpha = fg[:, :, 3:] / 255.0
         bg[y:y+h, x:x+w] = alpha * fg[:, :, :3] + (1 - alpha) * bg[y:y+h, x:x+w]
     else:
-        bg[y:y+h, x:x+w] = fg[:, :, :3]
+        bg[y:y+h, x:x+w] = fg
     return bg
 
-# 壓縮成 zip
-def zip_output_files():
-    with zipfile.ZipFile(ZIP_PATH, "w", zipfile.ZIP_DEFLATED) as zipf:
-        for file in sorted(os.listdir(OUTPUT_DIR)):
-            if file.endswith(".jpg"):
-                zipf.write(os.path.join(OUTPUT_DIR, file), arcname=file)
+# 壓縮 ZIP
+def zip_results():
+    with zipfile.ZipFile(ZIP_PATH, "w", zipfile.ZIP_DEFLATED) as z:
+        for f in sorted(os.listdir(OUTPUT_DIR)):
+            if f.endswith(".jpg"):
+                z.write(os.path.join(OUTPUT_DIR, f), arcname=f)
     return ZIP_PATH
 
-# 主流程
-def process_pdf(pdf_file, prefix="result_"):
+# 合成流程
+def process_and_combine(fg_paths, prefix):
+    bg = cv2.imread(BG_PATH)
+    if bg is None:
+        raise gr.Error("❌ 無法讀取 background.jpg")
+
+    # 清空舊檔
+    for f in os.listdir(OUTPUT_DIR):
+        if f.endswith(".jpg"):
+            os.remove(os.path.join(OUTPUT_DIR, f))
+
+    results = []
+    for i, path in enumerate(fg_paths, 1):
+        fg = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+        if fg is None:
+            continue
+        fg = cv2.resize(fg, (TARGET_WIDTH, TARGET_HEIGHT))
+        result = overlay_image(bg.copy(), fg, *POS1)
+        result = overlay_image(result, fg, *POS2)
+
+        out_path = os.path.join(OUTPUT_DIR, f"{prefix}{i:03d}.jpg")
+        cv2.imwrite(out_path, result)
+        results.append((out_path, f"{prefix}{i:03d}"))
+
+        if i % 10 == 0:
+            gc.collect()
+
+    if not results:
+        raise gr.Error("⚠️ 沒有生成任何圖片")
+
+    zip_file = zip_results()
+    return results, zip_file, f"✅ 成功合成 {len(results)} 張圖片"
+
+# 主入口：支援圖片或 PDF 混合上傳
+def process_inputs(files, prefix):
     if not prefix.strip():
         prefix = "result_"
-    try:
-        progress = gr.Progress()
-        progress(0, desc="準備中...")
-
-        # 轉換PDF為圖片
-        progress(0.1, "正在轉換PDF頁面為圖片...")
-        image_paths = pdf_to_images(pdf_file)
-        
-        if len(image_paths) == 0:
-            return [], None, "PDF轉換失敗，未產生任何圖片"
-        if len(image_paths) > MAX_PAGES:
-            raise gr.Error(f"PDF頁數過多，最多支援 {MAX_PAGES} 頁")
-
-        # 讀取背景圖
-        bg = cv2.imread(BG_PATH)
-        if bg is None:
-            raise gr.Error("無法讀取背景圖（請確認 background.jpg 是否存在）")
-
-        # 清除舊結果
-        for f in os.listdir(OUTPUT_DIR):
-            if f.endswith(".jpg"):
-                os.remove(os.path.join(OUTPUT_DIR, f))
-
-        results = []
-        for i, img_path in enumerate(image_paths, 1):
-            progress(i / len(image_paths), f"合成第 {i} 頁...")
-
-            try:
-                fg = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
-                if fg is None:
-                    continue
 
-                fg = cv2.resize(fg, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_AREA)
-                result = overlay_image(bg.copy(), fg, *POS1)
-                result = overlay_image(result, fg, *POS2)
+    fg_paths = []
+    for f in files:
+        name = f.name.lower()
+        if name.endswith(".pdf"):
+            fg_paths.extend(pdf_to_images(f))
+        elif name.endswith((".png", ".jpg", ".jpeg")):
+            fg_paths.append(f.name)
+        else:
+            raise gr.Error(f"不支援的檔案格式：{name}")
 
-                output_path = f"{OUTPUT_DIR}/{prefix}{i:03d}.jpg"
-                cv2.imwrite(output_path, result)
-                results.append((output_path, f"Page {i:03d}"))
-
-                if i % 10 == 0:
-                    gc.collect()
-
-            except Exception as e:
-                print(f"處理第 {i} 頁失敗: {str(e)}")
-                continue
-
-        progress(1, "完成！")
-        zip_file = zip_output_files()
-        return results, zip_file, f"成功處理 {len(results)}/{len(image_paths)} 頁，已打包為 ZIP 可下載"
-
-    except Exception as e:
-        raise gr.Error(f"處理出錯：{str(e)}")
+    return process_and_combine(fg_paths, prefix)
 
 # Gradio 介面
 with gr.Blocks() as demo:
-    gr.Markdown("## 📄 PDF頁面合成工具（背景固定為 7872x1200）")
+    gr.Markdown("## 📄🖼️ PDF 或圖片合成工具（背景為 7872x1200）")
 
-    pdf_upload = gr.File(label="上傳PDF文件", file_types=[".pdf"])
-    prefix_input = gr.Textbox(label="輸出檔案前綴 (選填)", placeholder="例如：invoice_")
-    btn_run = gr.Button("開始合成", variant="primary")
+    with gr.Row():
+        upload_files = gr.Files(label="上傳 PDF 或圖片", file_types=[".pdf", ".png", ".jpg", ".jpeg"])
+        prefix_input = gr.Textbox(label="輸出檔案前綴 (選填)", placeholder="例如：project_")
+        btn_run = gr.Button("開始合成", variant="primary")
 
     with gr.Row():
-        gallery = gr.Gallery(label="合成結果預覽", columns=3, preview=True)
-        zip_download = gr.File(label="下載結果 ZIP 檔")
-        log_output = gr.Textbox(label="日誌", interactive=False)
+        gallery = gr.Gallery(label="預覽", columns=3)
+        zip_file = gr.File(label="下載 ZIP")
+        log_output = gr.Textbox(label="處理日誌", interactive=False)
 
     btn_run.click(
-    fn=process_pdf,
-    inputs=[pdf_upload, prefix_input],
-    outputs=[gallery, zip_download, log_output],
-    concurrency_limit=4
+        fn=process_inputs,
+        inputs=[upload_files, prefix_input],
+        outputs=[gallery, zip_file, log_output]
     )
 
 if __name__ == "__main__":