ai-swap/backend/examples/ai_swap_example.py

"""
AI换脸换装服务使用示例
展示如何使用AISwapService进行换脸换装操作
"""

import os
import sys
import time
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

from backend.services.ai_swap_service import AISwapService, process_swap_with_record
from backend.modules.database.operations import DatabaseOperations
from backend.utils.logger_config import setup_logger

logger = setup_logger(__name__)


def create_test_user():
    """创建测试用户"""
    db_ops = DatabaseOperations()
    
    # 检查用户是否已存在
    existing_user = db_ops.get_user_by_username("test_user")
    if existing_user:
        logger.info(f"测试用户已存在: {existing_user['id']}")
        return existing_user["id"]
    
    # 创建新用户
    user = db_ops.create_user(
        username="test_user",
        password_hash="test_password_hash",
        is_admin=False
    )
    
    logger.info(f"创建测试用户成功: {user['id']}")
    return user["id"]


def create_test_images(user_id: int):
    """创建测试图片记录"""
    db_ops = DatabaseOperations()
    
    # 检查是否已有测试图片
    existing_images = db_ops.get_user_images(user_id, page_size=10)
    if existing_images["total"] >= 2:
        logger.info("测试图片已存在")
        face_images = [img for img in existing_images["images"] if img["image_type"] == "face"]
        cloth_images = [img for img in existing_images["images"] if img["image_type"] == "cloth"]
        
        if face_images and cloth_images:
            return face_images[0]["id"], cloth_images[0]["id"]
    
    # 创建测试图片记录（这里只是创建记录，实际文件需要手动上传）
    # 在实际项目中，你需要先上传图片文件，然后创建记录
    
    # 模拟图片路径
    test_face_path = "backend/data/face.png"
    test_cloth_path = "backend/data/cloth.png"
    
    if not os.path.exists(test_face_path) or not os.path.exists(test_cloth_path):
        logger.warning("测试图片文件不存在，请先准备测试图片")
        return None, None
    
    # 创建人脸图片记录
    face_record = db_ops.create_image_record(
        user_id=user_id,
        image_type="face",
        original_filename="test_face.png",
        stored_path=test_face_path,
        file_size=os.path.getsize(test_face_path)
    )
    
    # 创建服装图片记录
    cloth_record = db_ops.create_image_record(
        user_id=user_id,
        image_type="cloth",
        original_filename="test_cloth.png",
        stored_path=test_cloth_path,
        file_size=os.path.getsize(test_cloth_path)
    )
    
    logger.info(f"创建测试图片记录成功: 人脸={face_record['id']}, 服装={cloth_record['id']}")
    return face_record["id"], cloth_record["id"]


def test_ai_swap_service():
    """测试AI换脸换装服务"""
    logger.info("开始测试AI换脸换装服务")
    
    try:
        # 1. 创建测试用户
        user_id = create_test_user()
        
        # 2. 创建测试图片记录
        face_image_id, cloth_image_id = create_test_images(user_id)
        
        if not face_image_id or not cloth_image_id:
            logger.error("无法获取测试图片，测试终止")
            return
        
        # 3. 创建服务实例
        ai_swap_service = AISwapService()
        
        # 4. 执行换脸换装
        logger.info("开始执行换脸换装...")
        start_time = time.time()
        
        result = ai_swap_service.process_swap_with_record(
            user_id=user_id,
            face_image_id=face_image_id,
            cloth_image_id=cloth_image_id,
            prompt="美女站在海边，阳光明媚，穿着时尚的连衣裙"
        )
        
        end_time = time.time()
        processing_time = end_time - start_time
        
        # 5. 输出结果
        if result["success"]:
            logger.info("✅ 换脸换装成功！")
            logger.info(f"处理时间: {processing_time:.2f}秒")
            logger.info(f"处理记录ID: {result['process_record_id']}")
            logger.info(f"结果图片ID: {result['result_image_id']}")
            logger.info(f"生成的文案: {result['copywriter_text']}")
            logger.info(f"历史提示词: {result['history_prompt']}")
            
            # 6. 查询历史记录
            logger.info("\n查询处理历史...")
            history = ai_swap_service.get_user_process_history(user_id, page=1, page_size=5)
            logger.info(f"用户共有 {history['total']} 条处理记录")
            
            # 7. 查询处理详情
            if history["records"]:
                latest_record = history["records"][0]
                detail = ai_swap_service.get_process_detail(latest_record["id"])
                if detail:
                    logger.info(f"最新处理详情: 记录ID={detail['process_record']['id']}")
                    logger.info(f"人脸图片: {detail['face_image']['original_filename']}")
                    logger.info(f"服装图片: {detail['cloth_image']['original_filename']}")
                    logger.info(f"结果图片: {detail['result_image']['original_filename']}")
            
        else:
            logger.error("❌ 换脸换装失败！")
            logger.error(f"错误信息: {result['error']}")
            logger.error(f"错误类型: {result['error_type']}")
    
    except Exception as e:
        logger.error(f"测试过程中发生异常: {str(e)}", exc_info=True)


def test_convenience_function():
    """测试便捷函数"""
    logger.info("开始测试便捷函数")
    
    try:
        # 使用便捷函数
        result = process_swap_with_record(
            user_id=1,
            face_image_id=1,
            cloth_image_id=2,
            prompt="美女在花园里拍照"
        )
        
        if result["success"]:
            logger.info("✅ 便捷函数测试成功！")
            logger.info(f"结果: {result}")
        else:
            logger.error("❌ 便捷函数测试失败！")
            logger.error(f"错误: {result['error']}")
    
    except Exception as e:
        logger.error(f"便捷函数测试异常: {str(e)}", exc_info=True)


def test_error_handling():
    """测试错误处理"""
    logger.info("开始测试错误处理")
    
    ai_swap_service = AISwapService()
    
    # 测试无效用户ID
    try:
        result = ai_swap_service.process_swap_with_record(
            user_id=99999,  # 不存在的用户ID
            face_image_id=1,
            cloth_image_id=2,
            prompt="测试"
        )
        logger.info(f"无效用户ID测试结果: {result}")
    except Exception as e:
        logger.info(f"无效用户ID测试异常: {str(e)}")
    
    # 测试无效图片ID
    try:
        result = ai_swap_service.process_swap_with_record(
            user_id=1,
            face_image_id=99999,  # 不存在的图片ID
            cloth_image_id=2,
            prompt="测试"
        )
        logger.info(f"无效图片ID测试结果: {result}")
    except Exception as e:
        logger.info(f"无效图片ID测试异常: {str(e)}")
    
    # 测试空提示词
    try:
        result = ai_swap_service.process_swap_with_record(
            user_id=1,
            face_image_id=1,
            cloth_image_id=2,
            prompt=""  # 空提示词
        )
        logger.info(f"空提示词测试结果: {result}")
    except Exception as e:
        logger.info(f"空提示词测试异常: {str(e)}")


def main():
    """主函数"""
    logger.info("=" * 50)
    logger.info("AI换脸换装服务测试开始")
    logger.info("=" * 50)
    
    # 1. 测试错误处理
    test_error_handling()
    
    # 2. 测试便捷函数
    test_convenience_function()
    
    # 3. 测试完整服务
    test_ai_swap_service()
    
    logger.info("=" * 50)
    logger.info("AI换脸换装服务测试完成")
    logger.info("=" * 50)


if __name__ == "__main__":
    main()