ai-swap/backend/modules/database

数据库模块使用说明

概述

数据库模块为AI换脸系统提供完整的数据存储和管理功能，包括：

• 数据库连接管理 - 支持SQLite和MySQL
• 数据模型定义 - 用户、图片、处理记录等表结构
• CRUD操作封装 - 完整的数据库操作接口
• 数据库迁移 - 版本管理和迁移功能

目录结构

database/
├── __init__.py              # 模块初始化
├── connection.py            # 数据库连接管理
├── models.py               # 数据模型定义
├── operations.py           # CRUD操作封装
├── migrations.py           # 数据库迁移
├── migrations/             # 迁移文件目录
│   └── example_migration.json
├── example_usage.py        # 使用示例
└── README.md              # 说明文档

快速开始

1. 安装依赖

pip install -r requirements_database.txt

2. 初始化数据库

from backend.modules.database import DatabaseMigration

# 初始化数据库
migration = DatabaseMigration()
if migration.init_database():
    print("数据库初始化成功")

3. 基本使用

from backend.modules.database import DatabaseOperations

# 创建数据库操作实例
db_ops = DatabaseOperations()

# 创建用户
user = db_ops.create_user(
    username="test_user",
    email="test@example.com"
)

# 创建图片记录
image_record = db_ops.create_image_record(
    user_id=user.id,
    image_type="face",
    original_filename="face.jpg",
    stored_path="/data/images/face_001.jpg"
)

# 创建处理记录
process_record = db_ops.create_process_record(
    user_id=user.id,
    face_image_id=image_record.id,
    cloth_image_id=cloth_image.id,
    prompt="美女站在海边"
)

数据模型

用户表 (users)

• id - 用户ID
• username - 用户名
• email - 邮箱
• api_key - API密钥
• is_active - 是否激活
• is_admin - 是否管理员

图片记录表 (image_records)

• id - 图片记录ID
• user_id - 用户ID
• image_type - 图片类型 (face/cloth/result)
• original_filename - 原始文件名
• stored_path - 存储路径
• file_size - 文件大小
• image_hash - 图片哈希值

AI换脸处理记录表 (process_records)

• id - 处理记录ID
• user_id - 用户ID
• face_image_id - 人脸图片ID
• cloth_image_id - 服装图片ID
• result_image_id - 生成结果图片ID
• prompt - 主要处理提示词
• generated_text - AI生成的文案内容
• text_style - 文案风格(正式/轻松/创意等)
• face_prompt - 人脸识别提示词
• cloth_prompt - 服装识别提示词
• seed - 随机种子
• workflow_version - 工作流版本
• model_version - AI模型版本
• status - 处理状态(pending/processing/completed/failed)
• progress - 处理进度(0-100)
• processing_time - 处理耗时(秒)
• gpu_usage - GPU使用率
• memory_usage - 内存使用量(MB)

系统配置表 (system_configs)

• id - 配置ID
• config_key - 配置键
• config_value - 配置值
• config_type - 配置类型
• is_public - 是否公开

主要功能

数据库连接管理

from backend.modules.database import DatabaseConnection

# SQLite连接
db_conn = DatabaseConnection(
    db_type="sqlite",
    database="backend/data/ai_swap.db"
)

# MySQL连接
db_conn = DatabaseConnection(
    db_type="mysql",
    host="localhost",
    port=3306,
    database="ai_swap",
    username="root",
    password="password"
)

# 测试连接
if db_conn.test_connection():
    print("连接成功")

用户管理

# 创建用户
user = db_ops.create_user("username", "email@example.com")

# 查询用户
user = db_ops.get_user_by_username("username")
user = db_ops.get_user_by_api_key("api_key")

# 更新用户
db_ops.update_user(user.id, is_active=False)

# 删除用户
db_ops.delete_user(user.id)

图片记录管理

# 创建图片记录
image = db_ops.create_image_record(
    user_id=user.id,
    image_type="face",
    original_filename="face.jpg",
    stored_path="/path/to/image.jpg"
)

# 查询用户图片
images = db_ops.get_user_images(user.id, image_type="face")

# 更新图片记录
db_ops.update_image_record(image.id, file_size=1024000)

# 删除图片记录（软删除）
db_ops.delete_image_record(image.id, soft_delete=True)

AI换脸处理记录管理

# 创建AI换脸处理记录
process = db_ops.create_process_record(
    user_id=user.id,
    face_image_id=face_image.id,
    cloth_image_id=cloth_image.id,
    prompt="时尚模特在巴黎街头",
    face_prompt="保持面部特征清晰",
    cloth_prompt="时尚服装风格",
    seed=42,
    workflow_version="2.0.0",
    model_version="stable-diffusion-xl"
)

# 保存结果图片和生成的文案
db_ops.save_result_image(
    process_id=process.id,
    result_image_id=result_image.id,
    generated_text="时尚模特在巴黎街头展现优雅气质",
    text_style="时尚"
)

# 更新性能指标
db_ops.update_processing_metrics(
    process_id=process.id,
    processing_time=15.5,
    gpu_usage=85.2,
    memory_usage=2048.0
)

# 获取用户AI换脸历史
history = db_ops.get_user_ai_swap_history(user.id, page=1, page_size=20)

# 获取成功的换脸记录
successful_swaps = db_ops.get_successful_swaps(user.id, limit=10)

# 搜索特定提示词的处理记录
search_results = db_ops.search_processes_by_prompt("海边", user.id)

系统配置管理

# 设置配置
db_ops.set_config(
    "max_image_size",
    10485760,
    config_type="int",
    description="最大图片文件大小"
)

# 获取配置值
max_size = db_ops.get_config_value("max_image_size", default=5242880)

# 获取配置列表
configs = db_ops.list_configs(is_public=True)

数据库迁移

from backend.modules.database import DatabaseMigration

migration = DatabaseMigration()

# 创建迁移文件
migration_file = migration.create_migration(
    version="1.1.0",
    description="添加新字段",
    up_sql="ALTER TABLE users ADD COLUMN new_field TEXT;",
    down_sql="ALTER TABLE users DROP COLUMN new_field;"
)

# 执行迁移
migration.migrate_up()

# 回滚迁移
migration.migrate_down("1.0.0")

# 检查数据库健康状态
health = migration.check_database_health()

AI换脸专用功能

完整工作流示例

from backend.modules.database import DatabaseOperations

db_ops = DatabaseOperations()

# 1. 创建用户
user = db_ops.create_user("ai_user", "ai@example.com")

# 2. 上传人脸图片
face_image = db_ops.create_image_record(
    user_id=user.id,
    image_type="face",
    original_filename="face.jpg",
    stored_path="/data/images/face_001.jpg"
)

# 3. 上传服装图片
cloth_image = db_ops.create_image_record(
    user_id=user.id,
    image_type="cloth",
    original_filename="cloth.jpg",
    stored_path="/data/images/cloth_001.jpg"
)

# 4. 创建处理记录
process = db_ops.create_process_record(
    user_id=user.id,
    face_image_id=face_image.id,
    cloth_image_id=cloth_image.id,
    prompt="美女站在海边",
    workflow_version="1.0.0",
    model_version="stable-diffusion-v1.5"
)

# 5. 处理完成后保存结果
result_image = db_ops.create_image_record(
    user_id=user.id,
    image_type="result",
    original_filename="result.jpg",
    stored_path="/data/images/result_001.jpg"
)

db_ops.save_result_image(
    process_id=process.id,
    result_image_id=result_image.id,
    generated_text="美丽的海边风景，阳光明媚",
    text_style="轻松"
)

性能监控

# 更新处理性能指标
db_ops.update_processing_metrics(
    process_id=process.id,
    processing_time=12.5,
    gpu_usage=85.2,
    memory_usage=2048.0
)

历史记录查询

# 获取用户AI换脸历史
history = db_ops.get_user_ai_swap_history(user.id, page=1, page_size=20)

# 获取成功的换脸记录
successful_swaps = db_ops.get_successful_swaps(user.id, limit=10)

# 搜索特定提示词
search_results = db_ops.search_processes_by_prompt("海边", user.id)

高级功能

统计信息

# 获取系统统计信息
stats = db_ops.get_statistics()
print(f"总用户数: {stats['total_users']}")
print(f"人脸图片数: {stats['face_images']}")
print(f"服装图片数: {stats['cloth_images']}")
print(f"结果图片数: {stats['result_images']}")
print(f"完成率: {stats['completion_rate']:.2f}%")
print(f"平均处理时间: {stats['avg_processing_time']}秒")

分页查询

# 分页获取用户列表
result = db_ops.list_users(page=1, page_size=20, is_active=True)
print(f"用户列表: {result['users']}")
print(f"总页数: {result['total_pages']}")

事务管理

# 使用上下文管理器自动管理事务
with db_connection.get_session() as session:
    # 执行数据库操作
    session.add(new_record)
    # 自动提交或回滚

配置说明

数据库配置

默认使用SQLite数据库，文件位置：backend/data/ai_swap.db

如需使用MySQL，请修改连接参数：

db_conn = DatabaseConnection(
    db_type="mysql",
    host="localhost",
    port=3306,
    database="ai_swap",
    username="your_username",
    password="your_password"
)

连接池配置

• pool_size: 连接池大小（默认10）
• max_overflow: 最大溢出连接数（默认20）
• pool_recycle: 连接回收时间（MySQL默认3600秒）

错误处理

所有数据库操作都包含完整的错误处理：

try:
    user = db_ops.create_user("username", "email")
except Exception as e:
    logger.error(f"创建用户失败: {e}")
    # 处理错误

性能优化

1. 索引优化: 所有常用查询字段都已建立索引
2. 连接池: 使用SQLAlchemy连接池管理数据库连接
3. 分页查询: 支持分页查询避免大量数据加载
4. 软删除: 图片记录支持软删除，避免数据丢失

安全考虑

1. SQL注入防护: 使用SQLAlchemy ORM防止SQL注入
2. 密码哈希: 用户密码使用哈希存储
3. API密钥: 每个用户都有唯一的API密钥
4. 权限控制: 支持管理员和普通用户权限区分

测试

运行示例代码测试数据库功能：

cd backend/modules/database
python example_usage.py

常见问题

Q: 如何切换数据库类型？

A: 修改DatabaseConnection的db_type参数，SQLite为"sqlite"，MySQL为"mysql"

Q: 如何处理数据库迁移？

A: 使用DatabaseMigration类创建和执行迁移文件

Q: 如何备份数据库？

A: SQLite直接复制数据库文件，MySQL使用mysqldump命令

Q: 如何优化查询性能？

A: 确保查询条件使用索引字段，使用分页查询避免大量数据加载