hnu202326010204
/
MuseGuard
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: 修改数据库初始化

Browse Source
pull/7/head
梁浩 5 months ago
parent 92156b90d0
commit 6c40a2f8f5
1 changed files with 49 additions and 42 deletions
Show Stats Download Patch File Download Diff File

91
src/backend/init_db.py
Unescape View File

@ -15,9 +15,9 @@ def init_database():
# 初始化角色数据
roles = [
{'role_id': 0, 'name': '管理员', 'max_concurrent_tasks': 15, 'description': '系统管理员，拥有最高权限'},
{'role_id': 1, 'name': 'VIP用户', 'max_concurrent_tasks': 10, 'description': '付费用户，享有较高的资源使用权限'},
{'role_id': 2, 'name': '普通用户', 'max_concurrent_tasks': 5, 'description': '免费用户，享有基本的资源使用权限'}
{'role_id': 0, 'role_code': 'admin', 'name': '管理员', 'max_concurrent_tasks': 15, 'description': '系统管理员，拥有最高权限'},
{'role_id': 1, 'role_code': 'vip', 'name': 'VIP用户', 'max_concurrent_tasks': 10, 'description': '付费用户，享有较高的资源使用权限'},
{'role_id': 2, 'role_code': 'normal', 'name': '普通用户', 'max_concurrent_tasks': 5, 'description': '免费用户，享有基本的资源使用权限'}
]
for role_data in roles:
existing = Role.query.filter_by(role_id=role_data['role_id']).first()
@ -26,29 +26,30 @@ def init_database():
db.session.add(new_role)
# 初始化任务状态数据
task_status = [
{'status_code': 'waiting', 'status_name': '待处理', 'description': '任务已创建，等待处理'},
{'status_code': 'processing', 'status_name': '进行中', 'description': '任务正在处理中'},
{'status_code': 'completed', 'status_name': '已完成', 'description':'任务已成功完成'},
{'status_code': 'failed', 'status_name': '失败', 'description': '任务处理失败'}
task_statuses = [
{'task_status_code': 'waiting', 'task_status_name': '待处理', 'description': '任务已创建，等待处理'},
{'task_status_code': 'processing', 'task_status_name': '进行中', 'description': '任务正在处理中'},
{'task_status_code': 'completed', 'task_status_name': '已完成', 'description':'任务已成功完成'},
{'task_status_code': 'failed', 'task_status_name': '失败', 'description': '任务处理失败'}
]
for status in task_status:
existing = TaskStatus.query.filter_by(status_code=status['status_code']).first()
for status in task_statuses:
existing = TaskStatus.query.filter_by(task_status_code=status['task_status_code']).first()
if not existing:
new_status = TaskStatus(**status)
db.session.add(new_status)
# 初始化图片类型数据
image_types = [
{'image_code': 'original', 'image_name': '原始图片', 'description': '用户上传的原始图像文件'},
{'image_code': 'perturbed', 'image_name': '加噪后图片', 'description': '经过扰动算法处理后的防护图像'},
{'image_code': 'original_generate', 'image_name': '原始图像生成图片', 'description': '利用原始图像训练模型后模型生成图片'},
{'image_code': 'perturbed_generate', 'image_name': '加噪后图像生成图片', 'description': '利用加噪后图像训练模型后模型生成图片'},
{'image_code': 'heatmap', 'image_name': '生成的热力图', 'description': '热力图'}
{'image_code': 'original', 'image_name': '原始图', 'description': '用户上传的原始图像'},
{'image_code': 'perturbed', 'image_name': '加噪图', 'description': '经过扰动算法处理后的防护图像'},
{'image_code': 'original_generate', 'image_name': '原始图像生成图', 'description': '使用原始图像训练后生成的图像'},
{'image_code': 'perturbed_generate', 'image_name': '加噪图像生成图', 'description': '使用加噪图像训练后生成的图像'},
{'image_code': 'heatmap', 'image_name': '热力图', 'description': '原始图与加噪图的差异热力图'},
{'image_code': 'report', 'image_name': '报告图', 'description': '任务评估指标可视化图表'}
]
for img_type in image_types:
existing = ImageType.query.filter_by(type_code=img_type['type_code']).first()
existing = ImageType.query.filter_by(image_code=img_type['image_code']).first()
if not existing:
new_type = ImageType(**img_type)
db.session.add(new_type)
@ -62,7 +63,7 @@ def init_database():
]
for config in perturbation_configs:
existing = PerturbationConfig.query.filter_by(method_code=config['method_code']).first()
existing = PerturbationConfig.query.filter_by(perturbation_code=config['perturbation_code']).first()
if not existing:
new_config = PerturbationConfig(**config)
db.session.add(new_config)
@ -75,47 +76,53 @@ def init_database():
]
for config in finetune_configs:
existing = FinetuneConfig.query.filter_by(method_code=config['method_code']).first()
existing = FinetuneConfig.query.filter_by(finetune_code=config['finetune_code']).first()
if not existing:
new_config = FinetuneConfig(**config)
db.session.add(new_config)
# 初始化数据集类型数据
dataset_types = [
{'data_type_id': 0, 'dataset_code': 'facial', 'dataset_name': '人脸数据集', 'description': '人脸类型的数据集'},
{'data_type_id': 1, 'dataset_code': 'art', 'dataset_name': '艺术品数据集', 'description': '艺术品类型的数据集'}
data_types = [
{'data_type_code': 'facial', 'data_type_prompt': 'a photo of sks person', 'description': '人脸类型的数据集'},
{'data_type_code': 'art', 'data_type_prompt': 'a painting in the style of sks', 'description': '艺术品类型的数据集'}
]
for dataset in dataset_types:
existing = DataType.query.filter_by(data_type_id=dataset['data_type_id']).first()
for data_type in data_types:
existing = DataType.query.filter_by(data_type_code=data_type['data_type_code']).first()
if not existing:
new_dataset = DataType(**dataset)
db.session.add(new_dataset)
new_data_type = DataType(**data_type)
db.session.add(new_data_type)
# 初始化任务类型数据
# 初始化任务类型数据（按执行逻辑顺序排列）
task_types = [
{'task_type_id': 0, 'task_code': 'perturbation', 'task_name': '加噪任务', 'description': '对图像进行加噪处理的任务'},
{'task_type_id': 1, 'task_code': 'finetune', 'task_name': '微调任务', 'description': '对模型进行微调训练的任务'},
{'task_type_id': 2, 'task_code': 'generation', 'task_name': '生成任务', 'description': '利用微调后模型进行图像生成的任务'}
{'task_type_id': 3, 'task_code': 'heatmap', 'task_name': '热力图任务', 'description': '计算X和X’的热力图的任务'}
{'task_type_code': 'perturbation', 'task_type_name': '加噪任务', 'description': '对图像进行扰动处理，生成防护图像'},
{'task_type_code': 'heatmap', 'task_type_name': '热力图任务', 'description': '可视化原始图与加噪图的差异热力图'},
{'task_type_code': 'finetune', 'task_type_name': '微调任务', 'description': '使用图像数据集对模型进行微调训练'},
{'task_type_code': 'evaluate', 'task_type_name': '评估任务', 'description': '评估微调后模型的生成效果和防护性能'}
]
for task_type in task_types:
existing = TaskType.query.filter_by(task_type_code=task_type['task_type_code']).first()
if not existing:
new_task_type = TaskType(**task_type)
db.session.add(new_task_type)
# 创建默认管理员用户
admin_users = [
{'username': 'admin1', 'email': 'admin1@museguard.com', 'role_id': 0},
{'username': 'admin2', 'email': 'admin2@museguard.com', 'role_id': 0},
{'username': 'admin3', 'email': 'admin3@museguard.com', 'role_id': 0}
# 创建默认测试用户（三种角色各一个）
test_users = [
{'username': 'admin_test', 'email': 'admin@test.com', 'password': 'admin123', 'role_id': 0},
{'username': 'vip_test', 'email': 'vip@test.com', 'password': 'vip123', 'role_id': 1},
{'username': 'normal_test', 'email': 'normal@test.com', 'password': 'normal123', 'role_id': 2}
]
for admin_data in admin_users:
existing = User.query.filter_by(username=admin_data['username']).first()
for user_data in test_users:
existing = User.query.filter_by(username=user_data['username']).first()
if not existing:
admin_user = User(**admin_data)
admin_user.set_password('admin123') # 默认密码
db.session.add(admin_user)
password = user_data.pop('password') # 取出密码
test_user = User(**user_data)
test_user.set_password(password)
db.session.add(test_user)
# 为管理员创建默认配置
# 为测试用户创建默认配置
db.session.flush() # 确保user.id可用
user_config = UserConfig(user_id=admin_user.id)
user_config = UserConfig(user_id=test_user.user_id)
db.session.add(user_config)
# 提交所有更改

Write Preview
Loading…
Cancel
Save