import subprocess
import os,time
from rembg import remove, new_session
from PIL import Image
import io
from docxtpl import DocxTemplate, InlineImage
from docx.shared import Inches
input_folder = ‘tupian’ # 輸入文件夾
kouchu_folder = ‘kouchu’ # 去背景圖像保存文件夾
output_folder = ‘dangan’ # 合成后的輸出文件夾
background_folder = ‘beijing’ # 背景文件夾
model_path = “mode/u2net.onnx” # 模型路徑
wendang_path = “wendang” # 模型路徑
pdf_path = “pdf” # 模型路徑
doc = DocxTemplate(“moban.docx”)
context = {
‘company_name’: “World company”,
}
如果輸出文件夾不存在,則創建它
if not os.path.exists(output_folder):
os.makedirs(output_folder)
獲取背景文件夾中的所有背景圖片
background_images = [f for f in os.listdir(background_folder) if f.lower().endswith((‘.png’, ‘.jpg’, ‘.jpeg’))]
Function to convert docx to PDF
def convert_to_pdf(docx_path, pdf_output_path):
subprocess.run([“/Applications/LibreOffice.app/Contents/MacOS/soffice”, “–headless”, “–convert-to”, “pdf”, docx_path, “–outdir”, pdf_output_path])
遍歷輸入文件夾中的所有圖片文件
for filename in os.listdir(input_folder):
input_path = os.path.join(input_folder, filename)
# 確保文件是圖片文件(你可以根據需要調整此條件)
if filename.lower().endswith(('.png', '.jpg', '.jpeg')):output_path = os.path.join(kouchu_folder, filename)# 讀取輸入圖片并去背景with open(input_path, 'rb') as input_file:input_data = input_file.read()session = new_session(model_name='u2net_custom', model_path=model_path)output_data = remove(input_data, session=session)# 將去除背景后的數據轉為圖片input_image = Image.open(io.BytesIO(output_data))with open(output_path, 'wb') as output_file:output_file.write(output_data)# 遍歷背景文件夾中的所有背景圖像for bg_filename in background_images:background_path = os.path.join(background_folder, bg_filename)background_image = Image.open(background_path)# 調整背景大小,使其與輸入圖像大小匹配background_image = background_image.resize(input_image.size)# 合成圖像:將去背景后的圖片合成到背景上background_image = background_image.convert("RGBA")input_image = input_image.convert("RGBA")# 合成圖像(將前景圖像疊加到背景上)combined_image = Image.alpha_composite(background_image, input_image)# 將合成后的圖像轉換為RGB模式并保存combined_image = combined_image.convert("RGB")# 創建輸出路徑,根據背景文件名生成不同的輸出文件名output_filename = f"{os.path.splitext(filename)[0]}_{os.path.splitext(bg_filename)[0]}.jpg"output_path = os.path.join(output_folder, output_filename)# 保存合成后的圖片combined_image.save(output_path)print(f"處理并保存了 {filename} 與 {bg_filename} 合成的圖片 到 {output_path}")context['company_logo'] = InlineImage(doc, output_path, width=Inches(2)) # 插入圖片并設置寬度print(context)# 渲染模板doc.render(context)# 保存生成的文檔doc_output_path = os.path.join(wendang_path, f"generated_doc_{os.path.splitext(filename)[0]}_{os.path.splitext(bg_filename)[0]}.docx")doc.save(doc_output_path)# 延遲2秒time.sleep(2)pdf_output_path = os.path.join(pdf_path,f"generated_doc_{os.path.splitext(filename)[0]}_{os.path.splitext(bg_filename)[0]}.pdf")convert_to_pdf(doc_output_path, pdf_output_path)