Update app.py

app.py

@@ -1,137 +1,127 @@
-import cv2
-import numpy as np
 import os
+import cv2
 import gradio as gr
+import numpy as np
 import zipfile
-import tempfile
+import fitz  # PyMuPDF
 import gc
-from pdf2image import convert_from_path
-from PIL import Image
 
-# 設定常數與目錄
-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)
 
-# 確保目錄存在
-for d in [FG_DIR, OUTPUT_DIR]:
-    os.makedirs(d, exist_ok=True)
+# 建立資料夾
+def ensure_dir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+    if not os.access(path, os.W_OK):
+        raise RuntimeError(f"❌ 沒有寫入權限：{path}")
 
-# 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')
+ensure_dir(FG_DIR)
+ensure_dir(OUTPUT_DIR)
 
+# 將 PDF 每頁轉為圖檔
+def pdf_to_images(pdf_file):
+    doc = fitz.open(pdf_file.name)
     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)
+    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 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:
         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_inputs(pdf_file, image_files, prefix):
-    try:
-        progress = gr.Progress()
-        fg_images = []
+# 合成流程
+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(FG_DIR):
-            os.remove(os.path.join(FG_DIR, f))
-        for f in os.listdir(OUTPUT_DIR):
+    # 清空舊檔
+    for f in os.listdir(OUTPUT_DIR):
+        if f.endswith(".jpg"):
             os.remove(os.path.join(OUTPUT_DIR, f))
 
-        # PDF
-        if pdf_file is not None:
-            progress(0.1, "轉換 PDF 為圖片...")
-            fg_images = pdf_to_images(pdf_file, prefix)
+    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_"
+
+    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}")
 
-        # 圖片
-        elif image_files is not None:
-            progress(0.1, "儲存上傳圖片...")
-            fg_images = save_uploaded_images(image_files, prefix)
+    return process_and_combine(fg_paths, 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
+# Gradio 介面
 with gr.Blocks() as demo:
-    gr.Markdown("## 📄 PDF 或圖片合成工具（固定背景圖）")
+    gr.Markdown("## 📄🖼️ PDF 或圖片合成工具（背景為 7872x1200）")
 
     with gr.Row():
-        pdf_file = gr.File(label="上傳 PDF（選一）", file_types=[".pdf"])
-        image_files = gr.Files(label="或上傳圖片（JPG/PNG 多張）", file_types=[".jpg", ".png"])
+        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")
 
-    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)
+    with gr.Row():
+        gallery = gr.Gallery(label="預覽", columns=3)
+        zip_file = gr.File(label="下載 ZIP")
+        log_output = gr.Textbox(label="處理日誌", interactive=False)
 
-    run_btn.click(
+    btn_run.click(
         fn=process_inputs,
-        inputs=[pdf_file, image_files, prefix_input],
-        outputs=[gallery, zip_file, log],
-        concurrency_limit=2
+        inputs=[upload_files, prefix_input],
+        outputs=[gallery, zip_file, log_output]
     )
 
 if __name__ == "__main__":