122 lines
4.5 KiB
Python
122 lines
4.5 KiB
Python
import os
|
||
import sys
|
||
import json
|
||
from pathlib import Path
|
||
|
||
# 添加项目根目录到Python路径
|
||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
from server.core.database import get_db
|
||
from server.models.algorithm import Algorithm
|
||
from models.base import Base
|
||
import os
|
||
from core.database import engine
|
||
|
||
# 创建数据库表
|
||
Base.metadata.create_all(bind=engine)
|
||
|
||
def init_algorithms():
|
||
"""初始化算法数据，与algorithm目录保持一致"""
|
||
|
||
# 算法配置映射
|
||
algorithm_configs = {
|
||
"berthing": {
|
||
"name": "靠泊检测",
|
||
"description": "检测船舶靠泊过程中的关键行为",
|
||
"version": "1.0.0",
|
||
"detection_classes": ["person", "ship"],
|
||
"input_size": "640x640",
|
||
"tags": ["靠泊", "船舶", "安全"]
|
||
},
|
||
"unberthing": {
|
||
"name": "离泊检测",
|
||
"description": "检测船舶离泊过程中的关键行为",
|
||
"version": "1.0.0",
|
||
"detection_classes": ["person", "ship"],
|
||
"input_size": "640x640",
|
||
"tags": ["离泊", "船舶", "安全"]
|
||
},
|
||
"board_ship": {
|
||
"name": "登轮检测",
|
||
"description": "检测人员登轮行为",
|
||
"version": "1.0.0",
|
||
"detection_classes": ["person"],
|
||
"input_size": "640x640",
|
||
"tags": ["登轮", "人员", "安全"]
|
||
},
|
||
"leave_ship": {
|
||
"name": "离轮检测",
|
||
"description": "检测人员离轮行为",
|
||
"version": "1.0.0",
|
||
"detection_classes": ["person"],
|
||
"input_size": "640x640",
|
||
"tags": ["离轮", "人员", "安全"]
|
||
},
|
||
"bullet_frame": {
|
||
"name": "弹窗检测",
|
||
"description": "检测系统弹窗和异常界面",
|
||
"version": "1.0.0",
|
||
"detection_classes": ["window", "alert"],
|
||
"input_size": "640x640",
|
||
"tags": ["弹窗", "界面", "异常"]
|
||
}
|
||
}
|
||
|
||
db = next(get_db())
|
||
|
||
try:
|
||
# 获取algorithm目录的绝对路径
|
||
algorithm_dir = Path(__file__).parent.parent / "algorithm"
|
||
|
||
for algorithm_folder, config in algorithm_configs.items():
|
||
folder_path = algorithm_dir / algorithm_folder
|
||
|
||
if folder_path.exists():
|
||
# 检查模型文件
|
||
weights_path = folder_path / "weights" / "best.pt"
|
||
model_path = f"algorithm/{algorithm_folder}/weights/best.pt" if weights_path.exists() else None
|
||
|
||
# 检查是否已存在该算法
|
||
existing_algorithm = db.query(Algorithm).filter(
|
||
Algorithm.name == config["name"]
|
||
).first()
|
||
|
||
if existing_algorithm:
|
||
# 更新现有算法
|
||
existing_algorithm.model_path = model_path
|
||
existing_algorithm.detection_classes = json.dumps(config["detection_classes"])
|
||
existing_algorithm.input_size = config["input_size"]
|
||
existing_algorithm.tags = json.dumps(config["tags"])
|
||
existing_algorithm.status = "active" if model_path else "inactive"
|
||
print(f"更新算法: {config['name']}")
|
||
else:
|
||
# 创建新算法
|
||
algorithm = Algorithm(
|
||
name=config["name"],
|
||
description=config["description"],
|
||
version=config["version"],
|
||
model_path=model_path,
|
||
detection_classes=json.dumps(config["detection_classes"]),
|
||
input_size=config["input_size"],
|
||
tags=json.dumps(config["tags"]),
|
||
status="active" if model_path else "inactive",
|
||
accuracy=0.0, # 默认值
|
||
inference_time=0.0, # 默认值
|
||
is_enabled=True,
|
||
creator="system"
|
||
)
|
||
db.add(algorithm)
|
||
print(f"创建算法: {config['name']}")
|
||
|
||
db.commit()
|
||
print("算法数据初始化完成!")
|
||
|
||
except Exception as e:
|
||
db.rollback()
|
||
print(f"初始化算法数据失败: {e}")
|
||
raise
|
||
finally:
|
||
db.close()
|
||
|
||
if __name__ == "__main__":
|
||
init_algorithms() |