Update app.py

app.py

@@ -1,13 +1,10 @@
-import os
-import cv2
-import gradio as gr
+import os, cv2, zipfile, gc, gradio as gr
 import numpy as np
-import zipfile
 import fitz  # PyMuPDF
-import gc
+from PIL import Image
+from io import BytesIO
 
-# 參數設定
-FG_DIR = "foregrounds"
+# 常數設定
 OUTPUT_DIR = "output"
 ZIP_PATH = os.path.join(OUTPUT_DIR, "results.zip")
 BG_PATH = "background.jpg"
@@ -15,119 +12,95 @@ TARGET_WIDTH, TARGET_HEIGHT = 2133, 1200
 POS1 = (1032, 0)
 POS2 = (4666, 0)
 
-# 建立資料夾
+# 檢查資料夾權限
 def ensure_dir(path):
-    if not os.path.exists(path):
-        os.makedirs(path)
+    os.makedirs(path, exist_ok=True)
     if not os.access(path, os.W_OK):
-        raise RuntimeError(f"❌ 沒有寫入權限：{path}")
-
-ensure_dir(FG_DIR)
+        raise RuntimeError(f"❌ 無法寫入：{path}")
 ensure_dir(OUTPUT_DIR)
 
-# 將 PDF 每頁轉為圖檔
+# PDF 轉圖像（記憶體內處理）
 def pdf_to_images(pdf_file):
-    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
+    images = []
+    with fitz.open(stream=pdf_file.read(), filetype="pdf") as doc:
+        for page in doc:
+            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            images.append(img)
+    return images
 
 # 圖像合成功能
-def overlay_image(bg, fg, x, y):
-    h, w = fg.shape[:2]
+def overlay_image(bg, fg_img, x, y):
+    fg = np.array(fg_img.resize((TARGET_WIDTH, TARGET_HEIGHT)))
     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]
+        alpha = fg[:, :, 3] / 255.0
+        for c in range(3):
+            bg[y:y+TARGET_HEIGHT, x:x+TARGET_WIDTH, c] = (
+                alpha * fg[:, :, c] + (1 - alpha) * bg[y:y+TARGET_HEIGHT, x:x+TARGET_WIDTH, c]
+            )
     else:
-        bg[y:y+h, x:x+w] = fg
+        bg[y:y+TARGET_HEIGHT, x:x+TARGET_WIDTH] = fg
     return bg
 
-# 壓縮 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_and_combine(fg_paths, prefix):
+# 合成 + 輸出 ZIP
+def process_and_combine(images, 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}"))
-
+    for i, img in enumerate(images, 1):
+        combined = overlay_image(bg.copy(), img, *POS1)
+        combined = overlay_image(combined, img, *POS2)
+        output_path = os.path.join(OUTPUT_DIR, f"{prefix}{i:03d}.jpg")
+        cv2.imwrite(output_path, combined)
+        results.append((output_path, f"{prefix}{i:03d}"))
         if i % 10 == 0:
             gc.collect()
 
     if not results:
-        raise gr.Error("⚠️ 沒有生成任何圖片")
+        raise gr.Error("⚠️ 沒有產生圖片")
 
-    zip_file = zip_results()
-    return results, zip_file, f"✅ 成功合成 {len(results)} 張圖片"
+    with zipfile.ZipFile(ZIP_PATH, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for f in sorted(os.listdir(OUTPUT_DIR)):
+            if f.endswith(".jpg"):
+                zipf.write(os.path.join(OUTPUT_DIR, f), arcname=f)
+    return results, ZIP_PATH, f"✅ 共合成 {len(results)} 張圖，已打包 ZIP"
 
-# 主入口：支援圖片或 PDF 混合上傳
+# 主流程
 def process_inputs(files, prefix):
-    if not prefix.strip():
-        prefix = "result_"
-
-    fg_paths = []
-    for f in files:
-        name = f.name.lower()
+    prefix = prefix.strip() or "result_"
+    images = []
+    for file in files:
+        name = file.name.lower()
         if name.endswith(".pdf"):
-            fg_paths.extend(pdf_to_images(f))
+            images.extend(pdf_to_images(file))
         elif name.endswith((".png", ".jpg", ".jpeg")):
-            fg_paths.append(f.name)
+            img = Image.open(file).convert("RGBA")
+            images.append(img)
         else:
             raise gr.Error(f"不支援的檔案格式：{name}")
+    return process_and_combine(images, prefix)
 
-    return process_and_combine(fg_paths, prefix)
-
-# Gradio 介面
-with gr.Blocks() as demo:
-    gr.Markdown("## 📄🖼️ PDF 或圖片合成工具（背景為 7872x1200）")
-
-    with gr.Row():
-        upload_files = gr.Files(label="上傳 PDF 或圖片", file_types=[".pdf", ".png", ".jpg", ".jpeg"])
-        gallery = gr.Gallery(label="預覽", columns=3)
-        
+# Gradio UI
+with gr.Blocks(title="PDF合成工具") as demo:
+    gr.Markdown("### 📄🖼️ PDF / 圖片雙輸入合成工具（背景固定為 7872x1200）")
 
     with gr.Row():
-        prefix_input = gr.Textbox(label="輸出檔案前綴 (選填)", placeholder="例如：project_")
-        zip_file = gr.File(label="下載 ZIP")
-        
+        files = gr.Files(label="📎 上傳 PDF 或圖片", file_types=[".pdf", ".png", ".jpg", ".jpeg"])
+        prefix = gr.Textbox(label="檔名前綴（可選）", placeholder="例如：poster_")
+        run_btn = gr.Button("🚀 開始合成", variant="primary")
 
     with gr.Row():
-        btn_run = gr.Button("開始合成", variant="primary")
-        log_output = gr.Textbox(label="處理日誌", interactive=False)
-        
+        gallery = gr.Gallery(label="🖼️ 預覽", columns=3, height="auto")
+        zip_file = gr.File(label="⬇️ 下載 ZIP")
+        log = gr.Textbox(label="📄 日誌", interactive=False)
 
-    btn_run.click(
-        fn=process_inputs,
-        inputs=[upload_files, prefix_input],
-        outputs=[gallery, zip_file, log_output]
-    )
+    run_btn.click(fn=process_inputs, inputs=[files, prefix], outputs=[gallery, zip_file, log])
 
 if __name__ == "__main__":
     demo.launch()