kg_agent/post_evaluator_v4_image_loader.py

"""
ImageUploader 新实现 - 使用PIL Image对象
参考demo,直接下载图片到内存并转为PIL Image,不上传文件
"""
import asyncio
import requests
from PIL import Image
import io
from typing import Optional, List, Any


class ImageUploader:
    """图片加载器 - 下载图片并转为PIL Image对象(参考demo,使用内联数据方式)"""

    @staticmethod
    async def upload_images(image_urls: List[str]) -> tuple[List[Any], List[str]]:
        """
        批量下载图片并转为PIL Image对象

        Args:
            image_urls: 图片URL列表

        Returns:
            (image_objects, []) - PIL Image对象列表和空列表(保持接口兼容)
        """
        if not image_urls:
            return [], []

        print(f"      📥 准备加载 {len(image_urls)} 张图片(PIL Image方式)...")

        # 并发下载所有图片
        tasks = [ImageUploader._load_single_image(url, idx) for idx, url in enumerate(image_urls)]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # 分离成功和失败的结果
        image_objects = []

        for idx, result in enumerate(results):
            if isinstance(result, Exception):
                print(f"      ⚠️  图片{idx}加载失败: {str(result)[:50]}")
            elif result is not None:
                image_objects.append(result)

        print(f"      ✅ 成功加载 {len(image_objects)}/{len(image_urls)} 张图片")
        return image_objects, []  # 返回空列表作为temp_paths,因为不需要清理

    @staticmethod
    async def _load_single_image(image_url: str, idx: int) -> Optional[Any]:
        """
        下载单张图片并转为PIL Image对象

        Args:
            image_url: 图片URL
            idx: 图片索引(用于日志)

        Returns:
            PIL Image对象
        """
        try:
            # 下载图片到内存
            loop = asyncio.get_event_loop()
            response = await loop.run_in_executor(
                None,
                lambda: requests.get(image_url, timeout=30)
            )
            response.raise_for_status()

            # 转换为PIL Image对象
            image = Image.open(io.BytesIO(response.content))

            # 转换为RGB模式(Gemini推荐)
            if image.mode != 'RGB':
                image = image.convert('RGB')

            file_size_kb = len(response.content) / 1024
            print(f"      ✓ 图片{idx}加载成功 ({file_size_kb:.1f}KB, {image.size[0]}x{image.size[1]})")

            return image

        except Exception as e:
            print(f"      ✗ 图片{idx}加载失败: {str(e)[:60]}")
            return None