Update app.py

app.py

@@ -2,14 +2,13 @@ import cv2
 import numpy as np
 import os
 import gradio as gr
-import gc
 import zipfile
+import tempfile
+import gc
 from pdf2image import convert_from_path
 from PIL import Image
-import tempfile
 
-# 常量設定
-MAX_PAGES = 250
+# 設定常數與目錄
 BG_PATH = "background.jpg"
 FG_DIR = "foregrounds"
 OUTPUT_DIR = "output"
@@ -18,27 +17,36 @@ TARGET_WIDTH, TARGET_HEIGHT = 2133, 1200
 POS1 = (1032, 0)
 POS2 = (4666, 0)
 
-def check_write_permission(path):
-    if not os.path.exists(path):
-        os.makedirs(path, exist_ok=True)
-    elif not os.access(path, os.W_OK):
-        raise RuntimeError(f"❌ 目錄 '{path}' 存在但無法寫入")
-
-check_write_permission(FG_DIR)
-check_write_permission(OUTPUT_DIR)
-
-def pdf_to_images(pdf_file, prefix):
-    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp_pdf:
-        tmp_pdf.write(pdf_file.read())  # ← 這裡改成 .read()，支援 NamedString
-        tmp_pdf.flush()
-        images = convert_from_path(tmp_pdf.name, fmt='jpg')
-        paths = []
-        for i, img in enumerate(images, 1):
-            path = os.path.join(FG_DIR, f"{prefix}page_{i:03d}.jpg")
-            img.save(path, 'JPEG')
-            paths.append(path)
-        return paths
-
+# 確保目錄存在
+for d in [FG_DIR, OUTPUT_DIR]:
+    os.makedirs(d, exist_ok=True)
+
+# PDF 轉圖片
+def pdf_to_images(pdf_file_path, prefix):
+    with open(pdf_file_path, "rb") as f:
+        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp_pdf:
+            tmp_pdf.write(f.read())
+            tmp_pdf.flush()
+            images = convert_from_path(tmp_pdf.name, fmt='jpg')
+
+    paths = []
+    for i, img in enumerate(images, 1):
+        path = os.path.join(FG_DIR, f"{prefix}page_{i:03d}.jpg")
+        img.save(path, 'JPEG')
+        paths.append(path)
+    return paths
+
+# 儲存上傳圖片
+def save_uploaded_images(files, prefix):
+    paths = []
+    for i, file in enumerate(files, 1):
+        img = Image.open(file.name).convert("RGB")
+        path = os.path.join(FG_DIR, f"{prefix}img_{i:03d}.jpg")
+        img.save(path, "JPEG")
+        paths.append(path)
+    return paths
+
+# 合成圖層
 def overlay_image(bg, fg, x, y):
     h, w = fg.shape[:2]
     if fg.shape[2] == 4:
@@ -48,81 +56,82 @@ def overlay_image(bg, fg, x, y):
         bg[y:y+h, x:x+w] = fg[:, :, :3]
     return bg
 
-def zip_output(prefix):
-    zip_path = ZIP_PATH.replace("results", prefix.rstrip("_"))
-    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
-        for f in sorted(os.listdir(OUTPUT_DIR)):
-            if f.endswith(".jpg"):
-                zf.write(os.path.join(OUTPUT_DIR, f), arcname=f)
-    return zip_path
-
-def process_inputs(files, prefix):
-    progress = gr.Progress()
-    prefix = prefix.strip() or "results_"
-    prefix = prefix if prefix.endswith("_") else prefix + "_"
-
-    # 清空輸出
-    for f in os.listdir(OUTPUT_DIR):
-        if f.endswith(".jpg"):
+# 壓縮輸出成 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)
+    return ZIP_PATH
+
+# 主處理邏輯
+def process_inputs(pdf_file, image_files, prefix):
+    try:
+        progress = gr.Progress()
+        fg_images = []
+
+        # 清除舊資料
+        for f in os.listdir(FG_DIR):
+            os.remove(os.path.join(FG_DIR, f))
+        for f in os.listdir(OUTPUT_DIR):
             os.remove(os.path.join(OUTPUT_DIR, f))
 
-    # 載入背景圖
-    bg = cv2.imread(BG_PATH)
-    if bg is None:
-        raise gr.Error("❌ 找不到背景圖 background.jpg")
-
-    all_fgs = []
-
-    for file in files:
-        name = file.name.lower()
-        if name.endswith(".pdf"):
-            imgs = pdf_to_images(file, prefix)
-            all_fgs.extend(imgs)
-        elif name.endswith((".png", ".jpg", ".jpeg")):
-            new_path = os.path.join(FG_DIR, f"{prefix}{os.path.basename(name)}")
-            with open(new_path, "wb") as out:
-                out.write(file.read_bytes())
-            all_fgs.append(new_path)
-
-    if not all_fgs:
-        return [], None, "⚠️ 沒有找到可處理的圖片"
-
-    results = []
-    for i, fg_path in enumerate(all_fgs, 1):
-        progress(i / len(all_fgs), f"處理第 {i} 張...")
-        fg = cv2.imread(fg_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()
-
-    zip_path = zip_output(prefix)
-    return results, zip_path, f"✅ 成功處理 {len(results)} 張圖，ZIP 已可下載"
-
-# Gradio UI (相容���版)
+        # PDF
+        if pdf_file is not None:
+            progress(0.1, "轉換 PDF 為圖片...")
+            fg_images = pdf_to_images(pdf_file, prefix)
+
+        # 圖片
+        elif image_files is not None:
+            progress(0.1, "儲存上傳圖片...")
+            fg_images = save_uploaded_images(image_files, prefix)
+
+        else:
+            raise gr.Error("請上傳 PDF 或圖片！")
+
+        # 背景圖
+        bg = cv2.imread(BG_PATH)
+        if bg is None:
+            raise gr.Error("找不到背景圖 background.jpg")
+
+        results = []
+        for i, img_path in enumerate(fg_images, 1):
+            progress(0.2 + 0.7 * i / len(fg_images), f"合成第 {i} 張圖片...")
+            fg = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+            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}result_{i:03d}.jpg")
+            cv2.imwrite(out_path, result)
+            results.append((out_path, f"{prefix}{i:03d}"))
+            if i % 10 == 0:
+                gc.collect()
+
+        zip_file = zip_output_files()
+        return results, zip_file, f"✅ 成功處理 {len(results)} 張圖片，ZIP 可下載"
+
+    except Exception as e:
+        raise gr.Error(f"處理出錯：{str(e)}")
+
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🖼️ PDF / 圖片合成工具（搭配背景圖）")
-    gr.Markdown("背景圖為固定的 `background.jpg`，前景可為 PDF 或圖片。")
-
-    files_input = gr.Files(label="📂 上傳 PDF 或圖片", file_types=[".pdf", ".jpg", ".png", ".jpeg"])
-    prefix_input = gr.Textbox(label="📎 自訂輸出檔案前綴", placeholder="例如 myproject_")
+    gr.Markdown("## 📄 PDF 或圖片合成工具（固定背景圖）")
 
-    run_btn = gr.Button("🚀 開始合成", variant="primary")
+    with gr.Row():
+        pdf_file = gr.File(label="上傳 PDF（選一）", file_types=[".pdf"])
+        image_files = gr.Files(label="或上傳圖片（JPG/PNG 多張）", file_types=[".jpg", ".png"])
 
-    gallery = gr.Gallery(label="合成預覽", columns=3, height="auto")
-    zip_file = gr.File(label="📦 下載結果 ZIP 檔")
-    log_output = gr.Textbox(label="📜 處理日誌", interactive=False)
+    prefix_input = gr.Textbox(label="自定義檔名前綴", value="results_")
+    run_btn = gr.Button("開始合成")
+    gallery = gr.Gallery(label="預覽", columns=3)
+    zip_file = gr.File(label="下載 ZIP")
+    log = gr.Textbox(label="處理日誌", interactive=False)
 
     run_btn.click(
         fn=process_inputs,
-        inputs=[files_input, prefix_input],
-        outputs=[gallery, zip_file, log_output]
+        inputs=[pdf_file, image_files, prefix_input],
+        outputs=[gallery, zip_file, log],
+        concurrency_limit=2
     )
 
 if __name__ == "__main__":