ForcePilot/backend/package/yuxi/utils/image_processor.py
Wenjie Zhang 410dd47c14 refactor: 将后端代码迁移至 backend 目录
- 将 server/, src/, scripts/, test/ 等目录移动到 backend/ 目录下
- 使用 git rename 保留文件历史记录
- 更新 docker-compose.yml 和 api.Dockerfile 配置

WIP: 项目结构重构进行中
2026-03-24 11:08:12 +08:00

208 lines
7.3 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
图片处理工具模块
支持图片的格式转换、压缩、缩略图生成等功能
"""
import base64
import io
from PIL import ExifTags, Image
from yuxi.utils import logger
class ImageProcessor:
"""图片处理类"""
# 支持的图片格式
SUPPORTED_FORMATS = {"JPEG", "PNG", "WebP", "GIF", "BMP"}
# 最大文件大小5MB
MAX_FILE_SIZE = 5 * 1024 * 1024
# 缩略图尺寸
THUMBNAIL_SIZE = (200, 200)
def process_image(self, image_data: bytes, original_filename: str = "") -> dict:
"""
处理上传的图片
Args:
image_data: 图片二进制数据
original_filename: 原始文件名
Returns:
dict: 包含处理结果的字典
"""
try:
# 验证图片格式
img_format, _ = self._validate_image_format(image_data)
if img_format not in self.SUPPORTED_FORMATS:
raise ValueError(f"不支持的图片格式: {img_format}")
# 加载图片
with Image.open(io.BytesIO(image_data)) as img:
# 处理EXIF方向信息
img = self._fix_image_orientation(img)
# 生成缩略图
thumbnail_data = self._generate_thumbnail(img)
# 压缩主图片(如果需要)
processed_data, final_format = self._compress_image(img, img_format)
# 转换为 base64
base64_data = base64.b64encode(processed_data).decode("utf-8")
base64_thumbnail = base64.b64encode(thumbnail_data).decode("utf-8")
# 获取图片信息
width, height = img.size
mime_type = f"image/{final_format.lower()}"
return {
"success": True,
"image_content": base64_data,
"thumbnail_content": base64_thumbnail,
"width": width,
"height": height,
"format": final_format,
"mime_type": mime_type,
"size_bytes": len(processed_data),
"original_filename": original_filename,
}
except Exception as e:
logger.error(f"图片处理失败: {str(e)}")
return {"success": False, "error": str(e)}
def _validate_image_format(self, image_data: bytes) -> tuple[str, str]:
"""验证图片格式并返回格式信息"""
try:
with Image.open(io.BytesIO(image_data)) as img:
return img.format, img.mode
except Exception as e:
raise ValueError(f"无效的图片格式: {str(e)}")
def _fix_image_orientation(self, img: Image.Image) -> Image.Image:
"""根据EXIF信息修正图片方向"""
try:
if hasattr(img, "_getexif"):
exif = img._getexif()
if exif is not None:
for tag, value in exif.items():
if tag in ExifTags.TAGS and ExifTags.TAGS[tag] == "Orientation":
if value == 3:
img = img.rotate(180, expand=True)
elif value == 6:
img = img.rotate(270, expand=True)
elif value == 8:
img = img.rotate(90, expand=True)
break
except Exception as e:
logger.warning(f"修正图片方向失败,使用原始方向: {str(e)}")
return img
def _generate_thumbnail(self, img: Image.Image) -> bytes:
"""生成缩略图"""
try:
# 创建副本以避免修改原图
thumbnail = img.copy()
# 转换为RGB模式处理RGBA等格式
if thumbnail.mode != "RGB":
thumbnail = thumbnail.convert("RGB")
# 生成缩略图,保持宽高比
thumbnail.thumbnail(self.THUMBNAIL_SIZE, Image.Resampling.LANCZOS)
# 转换为JPEG格式
with io.BytesIO() as output:
thumbnail.save(output, format="JPEG", quality=85, optimize=True)
return output.getvalue()
except Exception as e:
logger.error(f"生成缩略图失败: {str(e)}")
# 如果缩略图生成失败返回一个1x1的透明图片
with io.BytesIO() as output:
empty_img = Image.new("RGB", (1, 1), color="white")
empty_img.save(output, format="JPEG", quality=85)
return output.getvalue()
def _compress_image(self, img: Image.Image, original_format: str) -> tuple[bytes, str]:
"""
压缩图片,如果超过大小限制
Args:
img: PIL Image对象
original_format: 原始格式
Returns:
Tuple[bytes, str]: (压缩后的图片数据, 最终格式)
"""
# 创建副本
processed_img = img.copy()
# 转换为RGB模式如果需要
if processed_img.mode in ("RGBA", "LA", "P"):
processed_img = processed_img.convert("RGB")
# 尝试保持原始格式但优先使用JPEG更好的压缩
target_format = "JPEG" if original_format != "PNG" else "PNG"
# 初始质量设置
quality = 85
with io.BytesIO() as output:
# 第一次保存以检查大小
processed_img.save(output, format=target_format, quality=quality, optimize=True)
compressed_data = output.getvalue()
# 如果文件大小合适,直接返回
if len(compressed_data) <= self.MAX_FILE_SIZE:
return compressed_data, target_format
# 如果文件太大,逐步降低质量
while len(compressed_data) > self.MAX_FILE_SIZE and quality > 10:
quality -= 10
output.seek(0)
output.truncate(0)
processed_img.save(output, format=target_format, quality=quality, optimize=True)
compressed_data = output.getvalue()
# 如果质量降到最低仍然太大,尝试缩小尺寸
if len(compressed_data) > self.MAX_FILE_SIZE:
# 逐步缩小尺寸
scale_factor = 0.9
while len(compressed_data) > self.MAX_FILE_SIZE and scale_factor > 0.3:
new_width = int(processed_img.width * scale_factor)
new_height = int(processed_img.height * scale_factor)
resized_img = processed_img.resize((new_width, new_height), Image.Resampling.LANCZOS)
output.seek(0)
output.truncate(0)
resized_img.save(output, format=target_format, quality=85, optimize=True)
compressed_data = output.getvalue()
scale_factor -= 0.1
return compressed_data, target_format
# 全局实例
image_processor = ImageProcessor()
def process_uploaded_image(image_data: bytes, filename: str = "") -> dict:
"""
处理上传的图片(便捷函数)
Args:
image_data: 图片二进制数据
filename: 文件名
Returns:
dict: 处理结果
"""
return image_processor.process_image(image_data, filename)