feat: 新增车辆检测Paddle模型及相关服务,优化依赖与代码兼容性
1. 新增3套PaddlePaddle车辆检测相关模型文件 2. 新增车辆检测服务类与违停检测功能 3. 更新服务依赖并添加环境初始化脚本与文档 4. 修复YOLO检测tensor转换兼容问题 5. 新增PyTorch版本兼容性修复逻辑 6. 扩展模型服务支持Paddle模型加载
This commit is contained in:
@@ -1,10 +1,36 @@
|
||||
import os
|
||||
import logging
|
||||
from ultralytics import YOLO
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
# PyTorch 2.6 兼容性修复
|
||||
os.environ.setdefault('TORCH_DISABLE_TORCH_GRAPH_OPTIMIZER', '1')
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
import torch
|
||||
if hasattr(torch, 'serialization'):
|
||||
# 导入 ultralytics 相关类以供 add_safe_globals 使用
|
||||
from ultralytics.nn.tasks import DetectionModel
|
||||
|
||||
# 注册所有需要的类
|
||||
torch.serialization.add_safe_globals([
|
||||
DetectionModel, # ultralytics 模型类
|
||||
torch.nn.modules.container.Sequential, # torch 序列类
|
||||
torch.nn.modules.Conv2d, # 卷积层
|
||||
torch.nn.modules.batchnorm.BatchNorm2d, # 批归一化
|
||||
torch.nn.modules.activation.ReLU, # 激活函数
|
||||
torch.nn.modules.Linear, # 线性层
|
||||
torch.nn.modules.Dropout, # Dropout 层
|
||||
torch.nn.modules.Upsample, # 上采样
|
||||
torch.nn.modules.PixelShuffle, # 像素重排
|
||||
])
|
||||
logger.info("✅ 检测到 PyTorch 2.6+,应用 ultralytics 兼容性修复")
|
||||
except (ImportError, AttributeError, NameError) as e:
|
||||
logger.warning(f"⚠️ PyTorch 兼容性修复失败: {e}")
|
||||
|
||||
from ultralytics import YOLO
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
class ModelService:
|
||||
def __init__(self):
|
||||
self.models: Dict[str, Union[YOLO, object]] = {}
|
||||
@@ -65,6 +91,24 @@ class ModelService:
|
||||
'size': '6MB',
|
||||
'description': '基于YOLOv8的徘徊检测模型',
|
||||
'name': '徘徊检测'
|
||||
},
|
||||
'vehicle_detection': {
|
||||
'path': os.path.join(base_dir, 'models', 'vehicle_detection_paddle', 'mot_ppyoloe_l_36e_ppvehicle', 'model.pdmodel'),
|
||||
'type': 'paddle',
|
||||
'classes': ['vehicle'],
|
||||
'labels': {'vehicle': '车辆'},
|
||||
'size': '181MB',
|
||||
'description': '基于PaddlePaddle PP-YOLOE-l的车辆检测和跟踪模型',
|
||||
'name': '车辆检测 (Paddle)'
|
||||
},
|
||||
'illegal_parking_detection': {
|
||||
'path': os.path.join(base_dir, 'models', 'vehicle_detection_paddle', 'mot_ppyoloe_l_36e_ppvehicle', 'model.pdmodel'),
|
||||
'type': 'paddle',
|
||||
'classes': ['vehicle'],
|
||||
'labels': {'vehicle': '车辆'},
|
||||
'size': '200MB',
|
||||
'description': '基于PaddlePaddle PP-YOLOE-l的违停检测模型,支持车牌识别',
|
||||
'name': '违停检测 (Paddle)'
|
||||
}
|
||||
}
|
||||
|
||||
@@ -112,31 +156,44 @@ class ModelService:
|
||||
# 处理 PaddleDetection 模型
|
||||
if config['type'] == 'paddle':
|
||||
try:
|
||||
from .paddle_detection_service import SmokingDetectionModel
|
||||
logger.info(f"正在加载 PaddlePaddle Docker 服务: {model_id}")
|
||||
model = SmokingDetectionModel()
|
||||
if model_id == 'smoking_detection_paddle':
|
||||
from .paddle_detection_service import SmokingDetectionModel
|
||||
logger.info(f"正在加载 PaddlePaddle 抽烟检测服务: {model_id}")
|
||||
model = SmokingDetectionModel()
|
||||
elif model_id in ['vehicle_detection', 'illegal_parking_detection']:
|
||||
from .vehicle_detection_service import VehicleDetectionModel
|
||||
logger.info(f"正在加载 PaddlePaddle 车辆检测服务: {model_id}")
|
||||
model = VehicleDetectionModel()
|
||||
else:
|
||||
logger.error(f"未知的 Paddle 模型类型: {model_id}")
|
||||
return None
|
||||
|
||||
self.models[model_id] = model
|
||||
logger.info(f"PaddlePaddle Docker 服务加载成功: {model_id}")
|
||||
logger.info(f"PaddlePaddle 服务加载成功: {model_id}")
|
||||
return model
|
||||
except Exception as e:
|
||||
logger.error(f"PaddlePaddle Docker 服务加载失败: {model_id}, 错误: {e}")
|
||||
logger.error(f"PaddlePaddle 服务加载失败: {model_id}, 错误: {e}")
|
||||
return None
|
||||
|
||||
# 处理 YOLO 模型
|
||||
model_path = config['path']
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
logger.error(f"模型文件不存在: {model_path}")
|
||||
logger.warning(f"模型文件不存在: {model_path},跳过加载 {model_id}")
|
||||
return None
|
||||
|
||||
try:
|
||||
logger.info(f"正在加载 YOLO 模型: {model_id} from {model_path}")
|
||||
model = YOLO(model_path)
|
||||
|
||||
self.models[model_id] = model
|
||||
logger.info(f"YOLO 模型加载成功: {model_id}")
|
||||
return model
|
||||
except Exception as e:
|
||||
logger.error(f"YOLO 模型加载失败: {model_id}, 错误: {e}")
|
||||
logger.error(f"模型路径: {model_path}")
|
||||
import traceback
|
||||
logger.error(f"详细错误: {traceback.format_exc()}")
|
||||
return None
|
||||
|
||||
def get_model(self, model_id: str) -> Optional[Union[YOLO, object]]:
|
||||
|
||||
Reference in New Issue
Block a user