import os
import uuid
import zipfile
from datetime import datetime
from pathlib import Path
from typing import List, Sequence, Optional, Tuple, Type
from fastapi import UploadFile
from sqlalchemy import func
from sqlmodel import Session, select
from common.biz_exception import BizException
from common.string_utils import snake_to_camel
from entity.device_model_relation import DeviceModelRelation
from entity.model import AlgoModel, AlgoModelCreate, AlgoModelUpdate, AlgoModelInfo
from common.global_thread_pool import GlobalThreadPool
from common.consts import NotifyChangeType
class ModelService:
def __init__(self, db: Session):
self.db = db
self.__model_change_callbacks = [] # 用于存储回调函数
self.thread_pool = GlobalThreadPool()
def register_change_callback(self, callback):
"""注册设备变化回调函数"""
self.__model_change_callbacks.append(callback)
def notify_change(self, algo_model_id, change_type):
"""当设备发生变化时，调用回调通知变化"""
for callback in self.__model_change_callbacks:
self.thread_pool.executor.submit(callback, algo_model_id, change_type)
def get_model_list(self,
name: Optional[str] = None,
remark: Optional[str] = None,
) -> Sequence[AlgoModel]:
statement = self.model_query(name, remark)
results = self.db.exec(statement)
return results.all()
def get_model_page(self,
name: Optional[str] = None,
remark: Optional[str] = None,
offset: int = 0,
limit: int = 10
) -> Tuple[Sequence[AlgoModelInfo], int]:
statement = self.model_query(name, remark)
# 查询总记录数
total_statement = select(func.count()).select_from(statement.subquery())
total = self.db.exec(total_statement).one()
# 添加分页限制
statement = statement.offset(offset).limit(limit)
# 执行查询并返回结果
model_list = self.db.exec(statement)
model_info_list: List[AlgoModelInfo] = []
if model_list:
for model in model_list:
model_info_list.append(AlgoModelInfo(
**model.dict(),
usage_status="使用中" if self.get_model_usage(model.id) else "未使用"
))
return model_info_list, total # 返回分页数据和总数
def model_query(self, name, remark):
# 构建查询语句
statement = select(AlgoModel)
if name:
statement = statement.where(AlgoModel.name.like(f"%{name}%"))
if remark:
statement = statement.where(AlgoModel.remark.like(f"%{remark}%"))
return statement
def process_zip(self, file: UploadFile):
model_dir = Path('weights/')
model_handle_dir = Path('model_handler/')
model_dir.mkdir(parents=True, exist_ok=True)
model_handle_dir.mkdir(parents=True, exist_ok=True)
# 支持的模型文件扩展名
SUPPORTED_MODEL_EXTENSIONS = {".pt", ".onnx", ".engine"}
# 临时保存上传文件
temp_path = Path(f"temp_upload_{uuid.uuid4()}.zip")
with open(temp_path, "wb") as temp_file:
temp_file.write(file.file.read())
model_file_path = None
handle_file_path = None
try:
with zipfile.ZipFile(temp_path, 'r') as zip_ref:
# 获取压缩包文件列表
file_list = zip_ref.namelist()
# 检查是否有模型权重文件
model_file = next((f for f in file_list if Path(f).suffix in SUPPORTED_MODEL_EXTENSIONS), None)
if not model_file:
raise BizException(
status_code=400,
message=f"Model weight file ({', '.join(SUPPORTED_MODEL_EXTENSIONS)}) is required in the zip."
)
# 解压模型文件到模型目录
zip_ref.extract(model_file, model_dir)
model_file_path = model_dir / model_file
# 检查是否有可选的 Python 脚本
handle_file = next((f for f in file_list if f.endswith(".py")), None)
if handle_file:
zip_ref.extract(handle_file, model_handle_dir)
handle_file_path = model_handle_dir / handle_file
except zipfile.BadZipFile:
raise BizException(status_code=400, message="Invalid zip file.")
finally:
# 删除临时文件
temp_path.unlink()
return str(model_file_path), str(handle_file_path) if handle_file_path else None
def create_model(self, model_data: AlgoModelCreate, file: UploadFile):
self.process_model_file(file, model_data)
model = AlgoModel.model_validate(model_data)
model.create_time = datetime.now()
model.update_time = datetime.now()
self.db.add(model)
self.db.commit()
self.db.refresh(model)
return model
def process_model_file(self, file, model):
model_file_path, handle_file_path = self.process_zip(file)
model.path = model_file_path
if handle_file_path:
model.handle_task = snake_to_camel(os.path.splitext(os.path.basename(handle_file_path))[0])
else:
model.handle_task = 'BaseModelHandler'
def update_model(self, model_data: AlgoModelUpdate, file: UploadFile):
model = self.db.get(AlgoModel, model_data.id)
if not model:
return None
update_data = model_data.dict(exclude_unset=True)
for key, value in update_data.items():
setattr(model, key, value)
model.update_time = datetime.now()
if file:
self.process_model_file(file, model)
self.db.add(model)
self.db.commit()
self.db.refresh(model)
self.notify_change(model.id, NotifyChangeType.MODEL_UPDATE)
return model
def delete_model(self, model_id: int):
model = self.db.get(AlgoModel, model_id)
if not model:
return None
# 查询 device_model_relation 中是否存在启用的绑定关系
statement = (
select(DeviceModelRelation)
.where(DeviceModelRelation.algo_model_id == model_id)
.where(DeviceModelRelation.is_use == 1)
)
relation_in_use = self.db.exec(statement).first()
# 如果存在启用的绑定关系，提示无法删除
if relation_in_use:
raise BizException(message=f"模型 {model.name} 正在被设备使用，无法删除")
self.db.delete(model)
self.db.commit()
return model
def get_models_in_use(self) -> Sequence[AlgoModel]:
"""获取所有在 device_model_relation 表里有启用绑定关系的模型信息"""
statement = (
select(AlgoModel)
.join(DeviceModelRelation, DeviceModelRelation.algo_model_id == AlgoModel.id)
.where(DeviceModelRelation.is_use == 1)
.group_by(AlgoModel.id)
)
results = self.db.exec(statement).all()
return results
def get_model_usage(self, algo_model_id) -> bool:
statement = (
select(DeviceModelRelation)
.where(
DeviceModelRelation.is_use == 1,
DeviceModelRelation.algo_model_id == algo_model_id,
)
)
result = self.db.exec(statement).all()
return len(result) > 0
def get_model_by_id(self, model_id):
return self.db.get(AlgoModel, model_id)
