AI_design_agent/scr/utils.py

"""
工具函数模块

包含图片处理、JSON解析等通用工具函数
"""

import os
import json
from PIL import Image

def parse_json_output(ai_response_text):
    """
    尝试将AI的回复解析为JSON对象
    
    Args:
        ai_response_text: AI返回的字符串
    
    Returns:
        (success, data): 
            success: 是否解析成功
            data: 解析后的数据字典，如果失败则包含error字段
    """
    try:
        # 1. 清洗数据：有时候 AI 会不听话地加上 ```json ... ```，这里做个简单的清洗
        cleaned_text = ai_response_text.strip()
        if cleaned_text.startswith("```"):
            cleaned_text = cleaned_text.replace("```json", "").replace("```", "")
        
        # 2. 尝试解析
        data = json.loads(cleaned_text)
        
        # 3. 验证关键字段是否存在
        if "success" in data and "prompt" in data:
            return True, data
        else:
            return False, {"error": "JSON格式合法，但缺少 success 或 prompt 字段"}

    except json.JSONDecodeError as e:
        return False, {"error": f"JSON解析失败: {str(e)}", "raw_text": ai_response_text}


def resize_proportional(image, target_width, target_height):
    """
    按照目标尺寸等比例缩放图片，保持宽高比
    计算长和宽的比例，取较小值以确保图片完全适应目标尺寸
    
    Args:
        image: PIL Image 对象
        target_width: 目标宽度
        target_height: 目标高度
    
    Returns:
        缩放后的 PIL Image 对象
    """
    width, height = image.size
    
    # 计算宽度和高度分别的比例
    width_ratio = target_width / width
    height_ratio = target_height / height
    
    # 取较小的比例，确保图片不会超出目标尺寸，同时保持宽高比
    ratio = min(width_ratio, height_ratio)
    
    new_width = int(width * ratio)
    new_height = int(height * ratio)
    
    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)


def center_image_on_white_background(image, target_width, target_height):
    """
    将图片居中放在白底图上
    
    Args:
        image: PIL Image 对象（需要放置的图片）
        target_width: 目标宽度
        target_height: 目标高度
    
    Returns:
        居中放置在白底图上的 PIL Image 对象
    """
    # 创建白底图
    white_background = Image.new('RGB', (target_width, target_height), (255, 255, 255))
    
    # 计算居中位置
    img_width, img_height = image.size
    x_offset = (target_width - img_width) // 2
    y_offset = (target_height - img_height) // 2
    
    # 如果图片是 RGBA 模式，需要处理透明度
    if image.mode == 'RGBA':
        white_background.paste(image, (x_offset, y_offset), image)
    else:
        white_background.paste(image, (x_offset, y_offset))
    
    return white_background


def prepare_second_image(second_image_path, first_image_size):
    """
    准备第二张图片：生成白底图，按照第一张图的长宽比例等比例缩放并居中放置
    
    Args:
        second_image_path: 第二张图片路径
        first_image_size: 第一张图片的尺寸 (width, height)
    
    Returns:
        处理后的 PIL Image 对象
    """
    # 读取第二张图片
    second_img = Image.open(second_image_path)
    
    # 转换为 RGB 模式（如果不是的话）
    if second_img.mode != 'RGB' and second_img.mode != 'RGBA':
        second_img = second_img.convert('RGB')
    
    target_width, target_height = first_image_size
    
    # 按照第一张图的长宽比例等比例缩放第二张图（保持宽高比）
    resized_second = resize_proportional(second_img, target_width, target_height)
    
    # 居中放在白底图上
    final_second = center_image_on_white_background(resized_second, target_width, target_height)
    
    return final_second


def horizontal_concatenate_images(image1_path, image2_path, output_path=None):
    """
    将两张图片横向拼接，第二张图会按照第一张图的尺寸进行调整
    
    Args:
        image1_path: 第一张图片路径（作为尺寸参考）
        image2_path: 第二张图片路径
        output_path: 输出图片路径，如果为None则自动生成
    
    Returns:
        输出图片路径
    """
    # 读取第一张图片
    first_img = Image.open(image1_path)
    
    # 转换为 RGB 模式（如果不是的话）
    if first_img.mode != 'RGB' and first_img.mode != 'RGBA':
        first_img = first_img.convert('RGB')
    
    # 获取第一张图的尺寸
    first_width, first_height = first_img.size
    
    # 准备第二张图片
    second_img = prepare_second_image(image2_path, (first_width, first_height))
    
    # 横向拼接两张图片
    # 总宽度 = 第一张图宽度 + 第二张图宽度
    # 总高度 = 两张图中较高的那个
    total_width = first_width + second_img.width
    total_height = max(first_height, second_img.height)
    
    # 创建拼接后的图片
    concatenated = Image.new('RGB', (total_width, total_height), (255, 255, 255))
    
    # 粘贴第一张图（左侧）
    concatenated.paste(first_img, (0, 0))
    
    # 粘贴第二张图（右侧）
    concatenated.paste(second_img, (first_width, 0))
    draw = ImageDraw.Draw(concatenated)
    
    # A. 画一条红色的分割线
    line_width = max(2, int(total_width * 0.005)) # 根据图片大小动态调整线宽
    draw.line([(first_width, 0), (first_width, total_height)], fill="red", width=line_width)
    
    # B. 准备字体 (尝试加载系统字体，为了防止中文乱码，我们主要用英文标签，模型读英文完全没问题)
    # 字体大小动态设为图片宽度的 4%
    font_size = int(total_width * 0.04)
    try:
        # 尝试加载常见字体（Windows/Linux兼容性尝试）
        # 注意：如果你的环境是 Linux docker，可能需要指定具体的 ttf 路径，例如 /usr/share/fonts/...
        font = ImageFont.truetype("arial.ttf", font_size)
    except IOError:
        try:
            font = ImageFont.truetype("DejaVuSans.ttf", font_size)
        except IOError:
            # 如果都找不到，使用默认字体（通常很小，但总比报错好）
            print("⚠️ 未找到系统字体，使用默认字体，可能较小")
            font = ImageFont.load_default()

    # C. 定义标签文本 (用英文更稳妥，Qwen识别英文非常准)
    text_left = "REFERENCE (Original)"
    text_right = "GENERATED (Flat Lay)"
    
    # D. 绘制文字（带描边，防止背景色干扰）
    # 左侧文字位置
    draw.text((20, 20), text_left, fill="red", font=font, stroke_width=2, stroke_fill="white")
    
    # 右侧文字位置 (起始点是第一张图宽度 + 偏移量)
    draw.text((first_width + 20, 20), text_right, fill="red", font=font, stroke_width=2, stroke_fill="white")   
    # 生成输出路径
    if output_path is None:
        base_name1 = os.path.splitext(os.path.basename(image1_path))[0]
        base_name2 = os.path.splitext(os.path.basename(image2_path))[0]
        output_dir = r"D:\线稿图\temp"
        output_path = os.path.join(output_dir, f"{base_name1}_{base_name2}_concatenated.jpg")
    
    # 保存拼接后的图片
    concatenated.save(output_path, quality=95)
    
    print(f"✅ 图片拼接完成: {output_path}")
    print(f"   第一张图尺寸: {first_width}x{first_height}")
    print(f"   第二张图原始尺寸: {Image.open(image2_path).size}")
    print(f"   第二张图处理后尺寸: {second_img.size}")
    print(f"   拼接后尺寸: {total_width}x{total_height}")
    
    return output_path


if __name__ == "__main__":
    # 示例用法
    image1 = r"H:\data\线稿图\S1261A097.jpg"
    image2 = r"D:\线稿图\平面图2\S1261A097.jpg"
    
    if os.path.exists(image1) and os.path.exists(image2):
        horizontal_concatenate_images(image1, image2)
    else:
        print("请修改示例中的图片路径")