Python pyside6编写一个广告图片生成器

from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                                 QHBoxLayout, QPushButton, QLabel, QLineEdit, QFileDialog,
                                 QSpinBox, QColorDialog, QTextEdit, QSlider, QDialog, QCheckBox,
                                 QScrollArea)
from PySide6.QtCore import Qt
from PySide6.QtGui import QPixmap, QImage
from PIL import Image, ImageDraw, ImageFont
import sys
import os
import time
 
class PreviewDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("去背景预览")
        self.setModal(False)  # 非模态对话框
        layout = QVBoxLayout(self)
        
        # 预览标签
        self.preview_label = QLabel()
        self.preview_label.setFixedSize(300, 300)
        self.preview_label.setStyleSheet("QLabel { background-color: #808080; }")  # 灰色背景
        self.preview_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.preview_label)
        
        # 关闭按钮
        close_btn = QPushButton("关闭预览")
        close_btn.clicked.connect(self.close)
        layout.addWidget(close_btn)
 
class HelpDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("使用说明")
        self.setGeometry(150, 150, 600, 400)
        
        layout = QVBoxLayout(self)
        
        # 创建滚动区域
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        layout.addWidget(scroll)
        
        # 创建内容容器
        content = QWidget()
        scroll.setWidget(content)
        content_layout = QVBoxLayout(content)
        
        # 读取说明文件
        try:
            with open("readme.md", "r", encoding="utf-8") as f:
                help_text = f.read()
        except:
            help_text = "找不到说明文件"
        
        # 创建文本显示
        text_display = QTextEdit()
        text_display.setReadOnly(True)
        text_display.setMarkdown(help_text)
        content_layout.addWidget(text_display)
        
        # 关闭按钮
        close_btn = QPushButton("关闭")
        close_btn.clicked.connect(self.close)
        layout.addWidget(close_btn)
 
class AdGenerator(QMainWindow):
    def __init__(self):
        super().__init__()
        # 设置默认输出路径
        self.output_dir = "D:\\PIC"
        # 确保输出目录存在
        if not os.path.exists(self.output_dir):
            try:
                os.makedirs(self.output_dir)
            except Exception as e:
                print(f"创建输出目录失败: {str(e)}")
                self.output_dir = os.getcwd()  # 如果创建失败，使用当前目录
        
        self.initUI()
        self.product_image_path = None
        self.bg_color = (30, 30, 30)
        self.text_color = (255, 255, 255)
        self.threshold = 220  # 默认阈值
        self.preview_dialog = None
        self.current_image = None
        self.remove_bg_enabled = True  # 默认启用去背景
        self.ad_width = 620   # 默认宽度
        self.ad_height = 420  # 默认高度
        self.help_dialog = None
        
    def initUI(self):
        self.setWindowTitle('广告图片生成器')
        self.setGeometry(100, 100, 800, 600)
        
        # 创建中心部件和布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # 创建预览标签
        self.preview_label = QLabel()
        self.preview_label.setFixedSize(620, 420)
        self.preview_label.setStyleSheet("QLabel { background-color: #1e1e1e; }")
        self.preview_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.preview_label)
        
        # 创建控制面板
        controls_layout = QHBoxLayout()
        
        # 左侧控制面板
        left_panel = QVBoxLayout()
        
        # 添加尺寸设置
        size_layout = QHBoxLayout()
        size_label = QLabel("广告尺寸:")
        self.width_input = QSpinBox()
        self.width_input.setRange(300, 1920)  # 设置合理的范围
        self.width_input.setValue(620)
        self.width_input.setSuffix(" px")
        self.width_input.valueChanged.connect(self.size_changed)
        
        size_x_label = QLabel("×")
        
        self.height_input = QSpinBox()
        self.height_input.setRange(200, 1080)
        self.height_input.setValue(420)
        self.height_input.setSuffix(" px")
        self.height_input.valueChanged.connect(self.size_changed)
        
        size_layout.addWidget(size_label)
        size_layout.addWidget(self.width_input)
        size_layout.addWidget(size_x_label)
        size_layout.addWidget(self.height_input)
        left_panel.addLayout(size_layout)
        
        # 文本输入
        text_layout = QVBoxLayout()
        text_label = QLabel("广告文本:")
        self.text_input = QTextEdit()
        self.text_input.setPlaceholderText("输入广告文本")
        self.text_input.setMaximumHeight(100)
        text_layout.addWidget(text_label)
        text_layout.addWidget(self.text_input)
        left_panel.addLayout(text_layout)
        
        # 图片选择按钮
        self.select_image_btn = QPushButton("选择产品图片")
        self.select_image_btn.clicked.connect(self.select_image)
        left_panel.addWidget(self.select_image_btn)
        
        # 背景颜色选择
        self.bg_color_btn = QPushButton("选择背景颜色")
        self.bg_color_btn.clicked.connect(self.select_bg_color)
        left_panel.addWidget(self.bg_color_btn)
        
        # 文字颜色选择
        self.text_color_btn = QPushButton("选择文字颜色")
        self.text_color_btn.clicked.connect(self.select_text_color)
        left_panel.addWidget(self.text_color_btn)
        
        # 阈值选择
        threshold_layout = QHBoxLayout()
        
        # 添加去背景开关
        self.remove_bg_checkbox = QCheckBox("启用去背景")
        self.remove_bg_checkbox.setChecked(True)
        self.remove_bg_checkbox.stateChanged.connect(self.toggle_remove_bg)
        left_panel.addWidget(self.remove_bg_checkbox)
        
        threshold_label = QLabel("去背景阈值:")
        self.threshold_value_label = QLabel("220")
        self.threshold_slider = QSlider(Qt.Horizontal)
        self.threshold_slider.setMinimum(180)
        self.threshold_slider.setMaximum(250)
        self.threshold_slider.setValue(220)
        self.threshold_slider.valueChanged.connect(self.threshold_changed)
        
        # 添加预览按钮
        self.preview_btn = QPushButton("预览去背景")
        self.preview_btn.clicked.connect(self.show_preview)
        
        threshold_layout.addWidget(threshold_label)
        threshold_layout.addWidget(self.threshold_slider)
        threshold_layout.addWidget(self.threshold_value_label)
        threshold_layout.addWidget(self.preview_btn)
        left_panel.addLayout(threshold_layout)
        
        controls_layout.addLayout(left_panel)
        
        # 右侧控制面板
        right_panel = QVBoxLayout()
        
        # 添加输出路径设置
        output_layout = QHBoxLayout()
        output_label = QLabel("输出路径:")
        self.output_path_label = QLabel(self.output_dir)
        self.output_path_label.setStyleSheet("padding: 5px; background-color: #f0f0f0; border: 1px solid #ccc;")
        self.output_path_btn = QPushButton("选择路径")
        self.output_path_btn.clicked.connect(self.select_output_path)
        
        output_layout.addWidget(output_label)
        output_layout.addWidget(self.output_path_label, stretch=1)
        output_layout.addWidget(self.output_path_btn)
        right_panel.addLayout(output_layout)
        
        # 生成按钮
        self.generate_btn = QPushButton("生成广告图片")
        self.generate_btn.clicked.connect(self.generate_ad)
        right_panel.addWidget(self.generate_btn)
        
        # 保存按钮
        self.save_btn = QPushButton("保存图片")
        self.save_btn.clicked.connect(self.save_image)
        right_panel.addWidget(self.save_btn)
        
        # 添加帮助按钮
        self.help_btn = QPushButton("使用说明")
        self.help_btn.clicked.connect(self.show_help)
        right_panel.addWidget(self.help_btn)
        
        controls_layout.addLayout(right_panel)
        layout.addLayout(controls_layout)
    
    def select_image(self):
        file_name, _ = QFileDialog.getOpenFileName(
            self,
            "选择产品图片",
            "",
            "图片文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        if file_name:
            self.product_image_path = file_name
            self.preview_image()
    
    def select_bg_color(self):
        color = QColorDialog.getColor()
        if color.isValid():
            self.bg_color = (color.red(), color.green(), color.blue())
            self.preview_image()
    
    def select_text_color(self):
        color = QColorDialog.getColor()
        if color.isValid():
            self.text_color = (color.red(), color.green(), color.blue())
            self.preview_image()
    
    def threshold_changed(self):
        self.threshold = self.threshold_slider.value()
        self.threshold_value_label.setText(str(self.threshold))
        if self.preview_dialog and self.preview_dialog.isVisible():
            self.show_preview()
        if self.product_image_path:
            self.preview_image()
    
    def toggle_remove_bg(self, state):
        self.remove_bg_enabled = bool(state)
        # 更新阈值控件的启用状态
        self.threshold_slider.setEnabled(self.remove_bg_enabled)
        self.threshold_value_label.setEnabled(self.remove_bg_enabled)
        self.preview_btn.setEnabled(self.remove_bg_enabled)
        # 如果有图片，立即更新预览
        if self.product_image_path:
            self.preview_image()
    
    def remove_background(self, image):
        if not self.remove_bg_enabled:
            return image.convert('RGBA')
            
        # 转换图片为RGBA模式
        image = image.convert('RGBA')
        data = image.getdata()
        new_data = []
        
        for item in data:
            # 如果像素接近白色，则设置为透明
            if item[0] >= self.threshold and item[1] >= self.threshold and item[2] >= self.threshold:
                new_data.append((255, 255, 255, 0))
            else:
                new_data.append(item)
        
        image.putdata(new_data)
        return image
    
    def show_preview(self):
        if not self.product_image_path:
            return
            
        if not self.preview_dialog:
            self.preview_dialog = PreviewDialog(self)
        
        try:
            # 加载原图并去背景
            img = Image.open(self.product_image_path)
            img = self.remove_background(img)
            
            # 调整图片大小以适应预览窗口
            preview_size = 300
            img.thumbnail((preview_size, preview_size), Image.Resampling.LANCZOS)
            
            # 转换为QImage
            img_data = img.convert("RGBA").tobytes("raw", "RGBA")
            qimg = QImage(img_data, img.size[0], img.size[1], QImage.Format_RGBA8888)
            
            # 显示在预览窗口中
            pixmap = QPixmap.fromImage(qimg)
            self.preview_dialog.preview_label.setPixmap(pixmap)
            
            # 显示预览窗口
            self.preview_dialog.show()
            
        except Exception as e:
            print(f"预览出错: {str(e)}")
    
    def preview_image(self):
        self.current_image = self.generate_ad(preview=True)
        if self.current_image:
            pixmap = QPixmap.fromImage(self.current_image)
            self.preview_label.setPixmap(pixmap.scaled(
                self.preview_label.size(),
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            ))
    
    def generate_ad(self, preview=False):
        width, height = self.ad_width, self.ad_height  # 使用设置的尺寸
        image = Image.new('RGB', (width, height), self.bg_color)
        draw = ImageDraw.Draw(image)
        
        if self.product_image_path and os.path.exists(self.product_image_path):
            try:
                product_img = Image.open(self.product_image_path)
                product_img = self.remove_background(product_img)
                
                # 严格计算2/3宽度
                right_width = int(width * (2/3))  # 确保宽度正好是2/3
                product_height = int((right_width * product_img.height) / product_img.width)
                
                # 如果高度超出，按高度缩放，但保持宽度不超过2/3
                if product_height > height:
                    product_height = height
                    calc_width = int((product_height * product_img.width) / product_img.height)
                    right_width = min(calc_width, right_width)  # 取较小值确保不超过2/3
                
                product_img = product_img.resize((right_width, product_height))
                
                # 计算粘贴位置（右侧居中）
                paste_x = width - right_width
                paste_y = (height - product_height) // 2
                
                # 粘贴产品图片
                image.paste(product_img, (paste_x, paste_y), product_img)
            except Exception as e:
                print(f"处理图片时出错: {str(e)}")
        
        try:
            # 根据广告尺寸调整字体大小
            font_size = min(32, int(height / 13))  # 动态计算合适的字体大小
            font = ImageFont.truetype("msyh.ttc", font_size)
        except:
            font = ImageFont.load_default()
        
        # 文本区域（使用左侧1/3的空间）
        text_area_width = int(width * (1/3))  # 确保文本区域是1/3
        text = self.text_input.toPlainText()
        
        # 改进的文本换行处理
        lines = []
        current_line = ""
        
        # 按字符分割文本（考虑中文字符）
        for char in text:
            test_line = current_line + char
            text_bbox = draw.textbbox((0, 0), test_line, font=font)
            if text_bbox[2] - text_bbox[0] > text_area_width - 40:  # 留出左右边距
                if current_line:
                    lines.append(current_line)
                    current_line = char
                else:
                    lines.append(char)
                    current_line = ""
            else:
                current_line += char
        
        if current_line:
            lines.append(current_line)
        
        # 计算文本总高度
        line_height = font.size + 8
        total_text_height = len(lines) * line_height
        
        # 绘制文本（左侧居中）
        start_y = (height - total_text_height) // 2
        for i, line in enumerate(lines):
            text_bbox = draw.textbbox((0, 0), line, font=font)
            text_width = text_bbox[2] - text_bbox[0]
            x = 30  # 增加左边距
            y = start_y + i * line_height
            draw.text((x, y), line, font=font, fill=self.text_color)
        
        if preview:
            # 转换为QImage用于预览
            data = image.tobytes("raw", "RGB")
            qimage = QImage(data, image.size[0], image.size[1], QImage.Format_RGB888)
            return qimage
        else:
            # 生成文件名（使用时间戳）
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            output_path = os.path.join(self.output_dir, f"ad_{timestamp}.png")
            
            # 创建图片
            width, height = self.ad_width, self.ad_height
            image = Image.new('RGB', (width, height), self.bg_color)
            draw = ImageDraw.Draw(image)
            
            if self.product_image_path and os.path.exists(self.product_image_path):
                try:
                    product_img = Image.open(self.product_image_path)
                    product_img = self.remove_background(product_img)
                    
                    # 严格计算2/3宽度
                    right_width = int(width * (2/3))  # 确保宽度正好是2/3
                    product_height = int((right_width * product_img.height) / product_img.width)
                    
                    # 如果高度超出，按高度缩放，但保持宽度不超过2/3
                    if product_height > height:
                        product_height = height
                        calc_width = int((product_height * product_img.width) / product_img.height)
                        right_width = min(calc_width, right_width)  # 取较小值确保不超过2/3
                    
                    product_img = product_img.resize((right_width, product_height))
                    
                    # 计算粘贴位置（右侧居中）
                    paste_x = width - right_width
                    paste_y = (height - product_height) // 2
                    
                    # 粘贴产品图片
                    image.paste(product_img, (paste_x, paste_y), product_img)
                except Exception as e:
                    print(f"处理图片时出错: {str(e)}")
            
            try:
                # 根据广告尺寸调整字体大小
                font_size = min(32, int(height / 13))  # 动态计算合适的字体大小
                font = ImageFont.truetype("msyh.ttc", font_size)
            except:
                font = ImageFont.load_default()
            
            # 文本区域（使用左侧1/3的空间）
            text_area_width = int(width * (1/3))  # 确保文本区域是1/3
            text = self.text_input.toPlainText()
            
            # 改进的文本换行处理
            lines = []
            current_line = ""
            
            # 按字符分割文本（考虑中文字符）
            for char in text:
                test_line = current_line + char
                text_bbox = draw.textbbox((0, 0), test_line, font=font)
                if text_bbox[2] - text_bbox[0] > text_area_width - 40:  # 留出左右边距
                    if current_line:
                        lines.append(current_line)
                        current_line = char
                    else:
                        lines.append(char)
                        current_line = ""
                else:
                    current_line += char
            
            if current_line:
                lines.append(current_line)
            
            # 计算文本总高度
            line_height = font.size + 8
            total_text_height = len(lines) * line_height
            
            # 绘制文本（左侧居中）
            start_y = (height - total_text_height) // 2
            for i, line in enumerate(lines):
                text_bbox = draw.textbbox((0, 0), line, font=font)
                text_width = text_bbox[2] - text_bbox[0]
                x = 30  # 增加左边距
                y = start_y + i * line_height
                draw.text((x, y), line, font=font, fill=self.text_color)
            
            # 保存图片
            try:
                image.save(output_path)
                print(f"图片已保存到: {output_path}")
                return output_path
            except Exception as e:
                print(f"保存图片失败: {str(e)}")
                return None
    
    def save_image(self):
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        default_name = f"ad_{timestamp}.png"
        file_name, _ = QFileDialog.getSaveFileName(
            self,
            "保存广告图片",
            os.path.join(self.output_dir, default_name),
            "PNG图片 (*.png)"
        )
        if file_name:
            if self.generate_ad():
                # 复制生成的图片到选择的位置
                latest_file = os.path.join(self.output_dir, default_name)
                if os.path.exists(latest_file):
                    Image.open(latest_file).save(file_name)
                    print(f"图片已保存到: {file_name}")
                else:
                    print("找不到生成的图片文件")
    
    def size_changed(self):
        # 更新尺寸
        self.ad_width = self.width_input.value()
        self.ad_height = self.height_input.value()
        
        # 更新预览标签大小
        self.preview_label.setFixedSize(self.ad_width, self.ad_height)
        
        # 如果有图片，重新生成预览
        if self.product_image_path:
            self.preview_image()
    
    def select_output_path(self):
        dir_path = QFileDialog.getExistingDirectory(
            self,
            "选择输出目录",
            self.output_dir,
            QFileDialog.ShowDirsOnly
        )
        if dir_path:
            self.output_dir = dir_path
            self.output_path_label.setText(self.output_dir)
    
    def show_help(self):
        if not self.help_dialog:
            self.help_dialog = HelpDialog(self)
        self.help_dialog.show()
 
if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = AdGenerator()
    window.show()
    sys.exit(app.exec())