369 lines
15 KiB
Python
369 lines
15 KiB
Python
import os
|
||
import asyncio
|
||
import hashlib
|
||
from typing import Dict, List, Tuple, Optional, Any
|
||
from dataclasses import dataclass
|
||
from PIL import Image
|
||
import io
|
||
import base64
|
||
from .logger import get_module_logger
|
||
logger = get_module_logger('image_optimizer')
|
||
from asyncio import Semaphore
|
||
import aiofiles
|
||
|
||
@dataclass
|
||
class ImageInfo:
|
||
"""图片信息数据类"""
|
||
filepath: str
|
||
filename: str
|
||
xref: Optional[int] = None
|
||
size: Optional[int] = None
|
||
hash: Optional[str] = None
|
||
compressed: bool = False
|
||
loaded: bool = False
|
||
data: Optional[bytes] = None
|
||
base64_data: Optional[str] = None
|
||
|
||
class ImageOptimizer:
|
||
"""图片优化处理器"""
|
||
|
||
def __init__(self,
|
||
max_concurrent_uploads: int = 5,
|
||
max_concurrent_api_calls: int = 3,
|
||
max_image_size_mb: float = 5.0,
|
||
compression_quality: int = 85,
|
||
batch_size: int = 10,
|
||
upload_max_retries: int = 3,
|
||
upload_retry_delay: float = 1.0):
|
||
"""
|
||
初始化图片优化器
|
||
|
||
Args:
|
||
max_concurrent_uploads: 最大并发上传数
|
||
max_concurrent_api_calls: 最大并发API调用数
|
||
max_image_size_mb: 图片最大尺寸(MB)
|
||
compression_quality: 压缩质量(1-100)
|
||
batch_size: 批处理大小
|
||
upload_max_retries: 上传失败最大重试次数
|
||
upload_retry_delay: 重试基础延迟时间(秒)
|
||
"""
|
||
self.max_concurrent_uploads = max_concurrent_uploads
|
||
self.max_concurrent_api_calls = max_concurrent_api_calls
|
||
self.max_image_size_bytes = max_image_size_mb * 1024 * 1024
|
||
self.compression_quality = compression_quality
|
||
self.batch_size = batch_size
|
||
self.upload_max_retries = upload_max_retries
|
||
self.upload_retry_delay = upload_retry_delay
|
||
|
||
# 并发控制信号量
|
||
self.upload_semaphore = Semaphore(max_concurrent_uploads)
|
||
self.api_semaphore = Semaphore(max_concurrent_api_calls)
|
||
|
||
# 缓存
|
||
self.hash_cache: Dict[str, str] = {} # hash -> uploaded_url
|
||
self.image_cache: Dict[str, ImageInfo] = {} # filepath -> ImageInfo
|
||
|
||
async def calculate_image_hash(self, filepath: str) -> str:
|
||
"""异步计算图片文件哈希值"""
|
||
hash_md5 = hashlib.md5()
|
||
async with aiofiles.open(filepath, 'rb') as f:
|
||
while chunk := await f.read(8192):
|
||
hash_md5.update(chunk)
|
||
return hash_md5.hexdigest()
|
||
|
||
async def load_image_info(self, filepath: str, filename: str, xref: Optional[int] = None) -> ImageInfo:
|
||
"""延迟加载图片信息，不立即读取内容"""
|
||
if filepath in self.image_cache:
|
||
return self.image_cache[filepath]
|
||
|
||
try:
|
||
stat = os.stat(filepath)
|
||
info = ImageInfo(
|
||
filepath=filepath,
|
||
filename=filename,
|
||
xref=xref,
|
||
size=stat.st_size,
|
||
loaded=False
|
||
)
|
||
self.image_cache[filepath] = info
|
||
return info
|
||
except Exception as e:
|
||
logger.error(f"Failed to load image info {filepath}: {e}")
|
||
raise
|
||
|
||
async def load_image_data(self, image_info: ImageInfo) -> bytes:
|
||
"""按需加载图片数据"""
|
||
if image_info.loaded and image_info.data:
|
||
return image_info.data
|
||
|
||
async with aiofiles.open(image_info.filepath, 'rb') as f:
|
||
image_info.data = await f.read()
|
||
image_info.loaded = True
|
||
|
||
# 计算哈希
|
||
if not image_info.hash:
|
||
image_info.hash = hashlib.md5(image_info.data).hexdigest()
|
||
|
||
return image_info.data
|
||
|
||
async def compress_image(self, image_data: bytes, max_size: Optional[int] = None) -> bytes:
|
||
"""压缩图片"""
|
||
if max_size is None:
|
||
max_size = self.max_image_size_bytes
|
||
|
||
# 如果图片已经小于限制，直接返回
|
||
if len(image_data) <= max_size:
|
||
return image_data
|
||
|
||
try:
|
||
# 使用PIL压缩图片
|
||
img = Image.open(io.BytesIO(image_data))
|
||
|
||
# 转换RGBA为RGB（如果需要）
|
||
if img.mode == 'RGBA':
|
||
rgb_img = Image.new('RGB', img.size, (255, 255, 255))
|
||
rgb_img.paste(img, mask=img.split()[3])
|
||
img = rgb_img
|
||
|
||
# 计算缩放比例
|
||
current_size = len(image_data)
|
||
scale_factor = (max_size / current_size) ** 0.5
|
||
|
||
if scale_factor < 1:
|
||
new_width = int(img.width * scale_factor)
|
||
new_height = int(img.height * scale_factor)
|
||
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||
|
||
# 压缩保存
|
||
output = io.BytesIO()
|
||
img.save(output, format='JPEG', quality=self.compression_quality, optimize=True)
|
||
compressed_data = output.getvalue()
|
||
|
||
logger.info(f"Compressed image from {len(image_data)} to {len(compressed_data)} bytes")
|
||
return compressed_data
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to compress image: {e}")
|
||
return image_data
|
||
|
||
async def process_image_for_upload(self, image_info: ImageInfo) -> Tuple[bytes, str]:
|
||
"""处理图片准备上传，返回(图片数据, 哈希值)"""
|
||
# 检查是否已经上传过相同内容
|
||
if image_info.hash and image_info.hash in self.hash_cache:
|
||
return None, image_info.hash
|
||
|
||
# 加载图片数据
|
||
image_data = await self.load_image_data(image_info)
|
||
|
||
# 压缩图片
|
||
if image_info.size > self.max_image_size_bytes:
|
||
image_data = await self.compress_image(image_data)
|
||
image_info.compressed = True
|
||
image_info.data = image_data
|
||
# 重新计算压缩后的哈希
|
||
image_info.hash = hashlib.md5(image_data).hexdigest()
|
||
|
||
return image_data, image_info.hash
|
||
|
||
async def batch_upload_images(self,
|
||
images: List[ImageInfo],
|
||
upload_func,
|
||
upload_options) -> Dict[str, str]:
|
||
"""批量并发上传图片"""
|
||
results = {}
|
||
upload_tasks = []
|
||
|
||
for image_info in images:
|
||
# 创建上传任务，使用配置的重试参数
|
||
task = self._upload_single_image(
|
||
image_info,
|
||
upload_func,
|
||
upload_options,
|
||
max_retries=self.upload_max_retries,
|
||
retry_delay=self.upload_retry_delay
|
||
)
|
||
upload_tasks.append(task)
|
||
|
||
# 并发执行所有上传任务
|
||
upload_results = await asyncio.gather(*upload_tasks, return_exceptions=True)
|
||
|
||
# 处理结果
|
||
for image_info, result in zip(images, upload_results):
|
||
if isinstance(result, Exception):
|
||
logger.error(f"Failed to upload {image_info.filename}: {result}")
|
||
# 即使失败也要记录结果，使用空字符串
|
||
if image_info.xref is not None:
|
||
results[image_info.xref] = ''
|
||
else:
|
||
results[image_info.filename] = ''
|
||
continue
|
||
|
||
url, upload_key = result
|
||
if url:
|
||
# 保存到缓存
|
||
if image_info.hash:
|
||
self.hash_cache[image_info.hash] = url
|
||
|
||
if image_info.xref is not None:
|
||
results[image_info.xref] = url
|
||
else:
|
||
results[image_info.filename] = url
|
||
else:
|
||
# URL为空的情况
|
||
logger.warning(f"Upload returned empty URL for {image_info.filename}")
|
||
if image_info.xref is not None:
|
||
results[image_info.xref] = ''
|
||
else:
|
||
results[image_info.filename] = ''
|
||
|
||
return results
|
||
|
||
async def _upload_single_image(self,
|
||
image_info: ImageInfo,
|
||
upload_func,
|
||
upload_options,
|
||
max_retries: int = 3,
|
||
retry_delay: float = 1.0) -> Tuple[Optional[str], Optional[str]]:
|
||
"""上传单个图片（带并发控制和重试机制）"""
|
||
# 注释掉测试模式，让上传回调能够被调用
|
||
# if os.getenv('MINERU_TEST_FILE'):
|
||
# return image_info.filepath, None
|
||
async with self.upload_semaphore:
|
||
# 处理图片
|
||
image_data, hash_value = await self.process_image_for_upload(image_info)
|
||
|
||
# 如果已经上传过，直接返回URL
|
||
if image_data is None and hash_value in self.hash_cache:
|
||
url = self.hash_cache[hash_value]
|
||
logger.info(f"Image {image_info.filename} already uploaded, reusing URL")
|
||
return url, None
|
||
|
||
# 重试机制
|
||
for attempt in range(max_retries):
|
||
try:
|
||
# 调用上传函数
|
||
logger.info(f"Uploading image {image_info.filename} (attempt {attempt + 1}/{max_retries})")
|
||
result = await upload_func(
|
||
image_info.filepath if not image_info.compressed else None,
|
||
image_info.filename,
|
||
upload_options,
|
||
binary_data=image_data if image_info.compressed else None
|
||
)
|
||
|
||
# 检查上传结果
|
||
if result and result[0]: # URL不为空
|
||
logger.info(f"Successfully uploaded {image_info.filename}")
|
||
return result
|
||
else:
|
||
logger.warning(f"Upload returned empty result for {image_info.filename}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"Upload attempt {attempt + 1} failed for {image_info.filename}: {e}")
|
||
|
||
# 如果不是最后一次尝试，等待后重试
|
||
if attempt < max_retries - 1:
|
||
await asyncio.sleep(retry_delay * (attempt + 1)) # 指数退避
|
||
continue
|
||
else:
|
||
# 最后一次尝试失败，记录详细错误
|
||
logger.error(f"All upload attempts failed for {image_info.filename}: {e}")
|
||
raise
|
||
|
||
# 所有重试都失败
|
||
return None, None
|
||
|
||
async def prepare_images_for_api(self, images: List[ImageInfo]) -> List[Dict[str, Any]]:
|
||
"""准备图片数据用于API调用"""
|
||
prepared = []
|
||
|
||
for image_info in images:
|
||
# 延迟加载图片数据
|
||
image_data = await self.load_image_data(image_info)
|
||
|
||
# 如果需要，压缩图片
|
||
if image_info.size > self.max_image_size_bytes:
|
||
image_data = await self.compress_image(image_data)
|
||
|
||
# 转换为base64
|
||
if not image_info.base64_data:
|
||
image_info.base64_data = base64.b64encode(image_data).decode('utf-8')
|
||
|
||
prepared.append({
|
||
'filename': image_info.filename,
|
||
'xref': image_info.xref,
|
||
'base64': image_info.base64_data,
|
||
'image_info': image_info
|
||
})
|
||
|
||
return prepared
|
||
|
||
async def batch_classify_images(self,
|
||
images: List[ImageInfo],
|
||
classify_func,
|
||
vision_model,
|
||
temp_dir: str,
|
||
src_name: str) -> Dict[int, Dict]:
|
||
"""批量分类图片 - 修改为顺序执行以避免对多模态服务造成压力"""
|
||
results = {}
|
||
|
||
# 分批处理
|
||
for i in range(0, len(images), self.batch_size):
|
||
batch = images[i:i + self.batch_size]
|
||
|
||
# 准备批量数据
|
||
prepared_images = await self.prepare_images_for_api(batch)
|
||
|
||
# 顺序执行分类，而不是并发
|
||
for img_data in prepared_images:
|
||
try:
|
||
# 顺序处理每个图片
|
||
result = await self._classify_single_image(
|
||
img_data,
|
||
classify_func,
|
||
vision_model,
|
||
temp_dir,
|
||
src_name
|
||
)
|
||
|
||
if img_data['xref'] is not None:
|
||
results[img_data['xref']] = result
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to classify image {img_data['filename']}: {e}")
|
||
continue
|
||
|
||
return results
|
||
|
||
async def _classify_single_image(self,
|
||
img_data: Dict[str, Any],
|
||
classify_func,
|
||
vision_model,
|
||
temp_dir: str,
|
||
src_name: str) -> Dict:
|
||
"""分类单个图片（带并发控制）"""
|
||
async with self.api_semaphore:
|
||
try:
|
||
# 调用分类函数
|
||
# classify_func expects: (learn_type, image_filepath, temp_dir, src_name, hint)
|
||
return await classify_func(
|
||
vision_model, # This is actually learn_type
|
||
img_data['image_info'].filepath,
|
||
temp_dir,
|
||
src_name,
|
||
hint=f"xref={img_data['xref']}"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Failed to classify image: {e}")
|
||
raise
|
||
|
||
def clear_cache(self):
|
||
"""清理缓存"""
|
||
self.image_cache.clear()
|
||
# 不清理hash_cache，因为它可以跨文档复用
|
||
|
||
async def cleanup(self):
|
||
"""清理资源"""
|
||
# 清理内存中的图片数据
|
||
for image_info in self.image_cache.values():
|
||
image_info.data = None
|
||
image_info.base64_data = None |