修复JPG导出为PNG，PS无法打开的问题

柴进
Showing 1 changed file with 137 additions and 11 deletions
image_generator.py
--- a/image_generator.py
View file @32c7b83
+++ b/image_generator.py
View file @32c7b83
@@ -271,6 +271,71 @@ def get_app_data_path() -> Path:
        return Path.cwd() / "images"
+def save_png_with_validation(file_path: str, image_bytes: bytes) -> bool:
+    """使用Pillow验证并重写PNG文件，确保格式正确性
+    Args:
+        file_path: 保存文件路径
+        image_bytes: 图像字节数据
+    Returns:
+        bool: True表示使用Pillow成功处理，False表示回退到原始方法
+    """
+    try:
+        from PIL import Image
+        import io
+        # 使用Pillow打开图像数据
+        with Image.open(io.BytesIO(image_bytes)) as img:
+            # 检查是否实际是JPEG但伪装成PNG的情况
+            file_format = img.format
+            if file_format == 'JPEG':
+                logger = logging.getLogger(__name__)
+                logger.info(f"检测到伪装PNG的JPEG文件，实际格式: {file_format}")
+            # 根据文件路径的扩展名决定保存格式
+            save_format = 'PNG' if file_path.lower().endswith('.png') else 'JPEG'
+            # 如果检测到的格式与目标格式不符，进行真正的转换
+            if file_format and file_format != save_format:
+                logger = logging.getLogger(__name__)
+                logger.info(f"执行格式转换: {file_format} -> {save_format}")
+                # 如果目标格式是PNG，确保使用RGBA模式以支持透明度
+                if save_format == 'PNG':
+                    if img.mode not in ['RGBA', 'RGB', 'L']:
+                        if img.mode == 'P':
+                            img = img.convert('RGBA')
+                        elif img.mode == 'LA':
+                            img = img.convert('RGBA')
+                        else:
+                            img = img.convert('RGBA')
+                elif save_format == 'JPEG':
+                    if img.mode in ['RGBA', 'P']:
+                        img = img.convert('RGB')
+                    elif img.mode == 'L':
+                        img = img.convert('RGB')
+            # 保存图片为指定格式
+            img.save(file_path, save_format, optimize=True)
+        logger = logging.getLogger(__name__)
+        logger.info(f"图片格式验证成功: {file_path}, 保存格式: {save_format}")
+        return True
+    except ImportError:
+        # Pillow不可用
+        logger = logging.getLogger(__name__)
+        logger.warning("Pillow库不可用，使用原始保存方法")
+        return False
+    except Exception as e:
+        # Pillow处理失败，回退到原始方法
+        logger = logging.getLogger(__name__)
+        logger.warning(f"Pillow处理失败，使用原始保存方法: {e}")
+        return False
 class HistoryManager:
    """历史记录管理器"""
@@ -310,17 +375,23 @@ class HistoryManager:
        record_dir = self.base_path / timestamp
        record_dir.mkdir(exist_ok=True)
-        # 保存生成的图片
+        # 保存生成的图片 - 使用PNG格式验证
        generated_image_path = record_dir / "generated.png"
-        with open(generated_image_path, 'wb') as f:
+        if not save_png_with_validation(str(generated_image_path), image_bytes):
-            f.write(image_bytes)
+            # Pillow处理失败，回退到原始方法
+            with open(generated_image_path, 'wb') as f:
+                f.write(image_bytes)
+            self.logger.warning(f"使用原始保存方法: {generated_image_path}")
-        # 保存参考图片
+        # 保存参考图片 - 使用PNG格式验证
        reference_image_paths = []
        for i, ref_img_bytes in enumerate(reference_images):
            ref_path = record_dir / f"reference_{i + 1}.png"
-            with open(ref_path, 'wb') as f:
+            if not save_png_with_validation(str(ref_path), ref_img_bytes):
-                f.write(ref_img_bytes)
+                # Pillow处理失败，回退到原始方法
+                with open(ref_path, 'wb') as f:
+                    f.write(ref_img_bytes)
+                self.logger.warning(f"使用原始保存方法: {ref_path}")
            reference_image_paths.append(ref_path)
        # 保存元数据
@@ -1929,20 +2000,63 @@ class ImageGeneratorWindow(QMainWindow):
            return
        try:
-            # Create temporary file
+            # 获取处理后的PNG图片数据
+            processed_bytes = self.get_processed_image_bytes('.png')
+            # Create temporary file with processed data
            with tempfile.NamedTemporaryFile(delete=False, suffix='.png', mode='wb') as tmp_file:
-                tmp_file.write(self.generated_image_bytes)
+                tmp_file.write(processed_bytes)
                tmp_path = tmp_file.name
            # Open with system default viewer
            url = QUrl.fromLocalFile(tmp_path)
            QDesktopServices.openUrl(url)
-            self.status_label.setText("● 已用系统查看器打开")
+            processing_method = "格式转换后" if processed_bytes != self.generated_image_bytes else "原始数据"
+            self.status_label.setText(f"● 已用系统查看器打开 ({processing_method})")
            self.status_label.setStyleSheet("QLabel { color: #007AFF; }")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"无法打开系统图片查看器: {str(e)}")
+    def get_processed_image_bytes(self, file_extension='.png') -> bytes:
+        """获取处理后的图片字节数据"""
+        try:
+            from PIL import Image
+            import io
+            # 使用Pillow处理图像
+            with Image.open(io.BytesIO(self.generated_image_bytes)) as img:
+                file_format = img.format
+                save_format = 'PNG' if file_extension.lower().endswith('.png') else 'JPEG'
+                # 检测并执行格式转换
+                if file_format and file_format != save_format:
+                    self.logger.info(f"下载时格式转换: {file_format} -> {save_format}")
+                    # 模式转换
+                    if save_format == 'PNG':
+                        if img.mode not in ['RGBA', 'RGB', 'L']:
+                            if img.mode == 'P':
+                                img = img.convert('RGBA')
+                            elif img.mode == 'LA':
+                                img = img.convert('RGBA')
+                            else:
+                                img = img.convert('RGBA')
+                    elif save_format == 'JPEG':
+                        if img.mode in ['RGBA', 'P']:
+                            img = img.convert('RGB')
+                        elif img.mode == 'L':
+                            img = img.convert('RGB')
+                # 转换为字节数据
+                processed_bytes = io.BytesIO()
+                img.save(processed_bytes, save_format, optimize=True)
+                return processed_bytes.getvalue()
+        except Exception as e:
+            self.logger.warning(f"图片处理失败，使用原始数据: {e}")
+            return self.generated_image_bytes
    def download_image(self):
        """Download generated image"""
        if not self.generated_image_bytes:
@@ -1961,14 +2075,26 @@ class ImageGeneratorWindow(QMainWindow):
        if file_path:
            try:
+                # 获取处理后的图片数据
+                file_extension = Path(file_path).suffix
+                processed_bytes = self.get_processed_image_bytes(file_extension)
+                # 保存处理后的图片
                with open(file_path, 'wb') as f:
-                    f.write(self.generated_image_bytes)
+                    f.write(processed_bytes)
+                # 获取实际文件大小
+                import os
+                file_size = os.path.getsize(file_path)
+                processing_method = "格式转换" if processed_bytes != self.generated_image_bytes else "原始数据"
+                self.logger.info(f"图片下载完成 - 方法: {processing_method}, 文件: {file_path}, 大小: {file_size} 字节")
-                file_size = len(self.generated_image_bytes)
                QMessageBox.information(self, "成功", f"图片已保存到:\n{file_path}\n\n文件大小: {file_size:,} 字节")
                self.status_label.setText("● 图片已保存")
                self.status_label.setStyleSheet("QLabel { color: #34C759; }")
            except Exception as e:
+                self.logger.error(f"图片保存失败: {str(e)}")
                QMessageBox.critical(self, "错误", f"保存失败: {str(e)}")
    def refresh_history(self):