AI_design_agent/scr/check.py

import json
import os
import requests
from io import BytesIO
from google import genai
from google.genai import types
from datetime import datetime
from google.genai import Client
from google.genai import types 
GEMINI_BASE_URL = os.environ.get("GEMINI_API_BASE_URL", "https://api.openaius.com")
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY","sk-qnsfJw0vsAitlnrXcOeBYrLbTv9LXfsN1m3jIUfMJagan5IR")

_client: Client | None = None

def _get_client() -> Client:
    global _client
    if _client is None:
        _client = Client(
            api_key=GEMINI_API_KEY,
            http_options={"base_url": GEMINI_BASE_URL}
        )
    return _client



def get_image_mime_type(image_url, content=None):
    """
    根据URL或内容判断图片的MIME类型
    
    Args:
        image_url: 图片URL
        content: 图片内容（可选）
    
    Returns:
        MIME类型字符串
    """
    # 首先尝试从URL判断
    if image_url.endswith('.png'):
        return 'image/png'
    elif image_url.endswith('.jpg') or image_url.endswith('.jpeg'):
        return 'image/jpeg'
    elif image_url.endswith('.gif'):
        return 'image/gif'
    elif image_url.endswith('.webp'):
        return 'image/webp'
    
    # 如果URL无法判断，尝试从内容判断
    if content:
        # 检查文件头
        if content.startswith(b'\x89PNG'):
            return 'image/png'
        elif content.startswith(b'\xff\xd8\xff'):
            return 'image/jpeg'
        elif content.startswith(b'GIF'):
            return 'image/gif'
        elif content.startswith(b'RIFF') and b'WEBP' in content[:12]:
            return 'image/webp'
    
    # 默认返回jpeg
    return 'image/jpeg'


def download_image_from_url(image_url):
    """
    从URL下载图片并返回字节数据和MIME类型
    
    Args:
        image_url: 图片URL
    
    Returns:
        (图片字节数据, MIME类型)，如果失败返回(None, None)
    """
    try:
        response = requests.get(image_url, timeout=30)
        response.raise_for_status()
        content = response.content
        
        # 尝试从响应头获取Content-Type
        content_type = response.headers.get('Content-Type', '')
        if content_type and content_type.startswith('image/'):
            mime_type = content_type
        else:
            # 从URL或内容判断
            mime_type = get_image_mime_type(image_url, content)
        
        return content, mime_type
    except Exception as e:
        print(f"❌ 下载图片失败 {image_url}: {e}")
        return None, None


def process_image_pair_with_gemini(image1_url, image2_url, prompt, model="gemini-2.5-flash"):
    """
    使用Gemini处理成对的图片
    
    Args:
        image1_url: 第一张图片的URL
        image2_url: 第二张图片的URL
        prompt: 提示词
        model: 使用的模型名称
    
    Returns:
        Gemini的响应文本，如果失败返回None
    """
    try:
        # 下载第一张图片并获取MIME类型
        img1_bytes, img1_mime = download_image_from_url(image1_url)
        if not img1_bytes:
            return None
        
        # 下载第二张图片并获取MIME类型
        img2_bytes, img2_mime = download_image_from_url(image2_url)
        if not img2_bytes:
            return None
        
        # 创建请求内容 - 直接使用内联图片数据，避免使用 files.upload
        contents = [
            prompt,
            types.Part.from_bytes(
                data=img1_bytes,
                mime_type=img1_mime
            ),
            types.Part.from_bytes(
                data=img2_bytes,
                mime_type=img2_mime
            )
        ]
        
        # 调用Gemini API
        response = _get_client().models.generate_content(
            model=model,
            contents=contents
        )
        
        return response.text
        
    except Exception as e:
        print(f"❌ 处理图片对时出错: {e}")
        return None





if __name__ == "__main__":
 

    prompt = """图2是不是图1的衣服的线稿款式图，而且除了衣服的线稿外没有其他的元素了？是不是只有黑白线稿没有颜色的，是不是不含人台（mannequin）的线稿，是不是衣服的数量和特征都对应上了，是的话回答是，否则回答否"""

    result=process_image_pair_with_gemini(    "https://testdgxcx-oss.gloria.com.cn/design_image/cropped_20251121_120610_072.png","https://testdgxcx-oss.gloria.com.cn/design_image/cropped_20251121_120624_088.png",prompt)
    print(result)