113 lines
3.6 KiB
Python
113 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
视频检测系统配置文件
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
|
|
class DetectionConfig:
|
|
"""检测系统配置类"""
|
|
|
|
# 模型配置
|
|
MODEL_PATH = "models/best.pt"
|
|
CONFIDENCE_THRESHOLD = 0.5
|
|
|
|
# 动态模型信息
|
|
_model_classes = None
|
|
_model_colors = None
|
|
|
|
# 文件夹配置
|
|
INPUT_FOLDER = "input_videos"
|
|
OUTPUT_FOLDER = "output_frames"
|
|
LOG_FOLDER = "logs"
|
|
|
|
# 视频处理配置
|
|
SUPPORTED_VIDEO_FORMATS = {'.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv', '.webm'}
|
|
PROGRESS_INTERVAL = 100 # 每处理多少帧显示一次进度
|
|
MIN_FRAME_INTERVAL = 1.0 # 检测帧的最小时间间隔（秒）
|
|
|
|
# 输出配置
|
|
SAVE_ORIGINAL_FRAMES = False # 是否保存原始帧
|
|
SAVE_ANNOTATED_FRAMES = True # 是否保存标注帧
|
|
SAVE_DETECTION_JSON = True # 是否保存检测结果JSON
|
|
|
|
# 图像质量配置
|
|
IMAGE_QUALITY = 95 # JPEG质量 (1-100)
|
|
|
|
# 默认道路损伤类别映射 (中文标签) - 作为备用
|
|
DEFAULT_CLASS_NAMES_CN = {
|
|
0: "纵向裂缝",
|
|
1: "横向裂缝",
|
|
2: "网状裂缝",
|
|
3: "坑洞",
|
|
4: "白线模糊"
|
|
}
|
|
|
|
# 默认类别颜色配置 (BGR格式) - 所有类别使用相同的绿色
|
|
DEFAULT_CLASS_COLORS = {
|
|
0: (0, 255, 0), # 绿色
|
|
1: (0, 255, 0), # 绿色
|
|
2: (0, 255, 0), # 绿色
|
|
3: (0, 255, 0), # 绿色
|
|
4: (0, 255, 0) # 绿色
|
|
}
|
|
|
|
@classmethod
|
|
def get_model_path(cls):
|
|
"""获取模型文件的绝对路径"""
|
|
current_dir = Path(__file__).parent
|
|
model_path = current_dir / cls.MODEL_PATH
|
|
return str(model_path.resolve())
|
|
|
|
@classmethod
|
|
def create_folders(cls):
|
|
"""创建必要的文件夹"""
|
|
folders = [cls.INPUT_FOLDER, cls.OUTPUT_FOLDER, cls.LOG_FOLDER]
|
|
for folder in folders:
|
|
Path(folder).mkdir(parents=True, exist_ok=True)
|
|
|
|
@classmethod
|
|
def load_model_classes(cls, model_path):
|
|
"""从模型文件动态加载类别信息"""
|
|
try:
|
|
model = YOLO(model_path)
|
|
cls._model_classes = model.names
|
|
# 为所有类别生成绿色
|
|
cls._model_colors = {i: (0, 255, 0) for i in range(len(model.names))}
|
|
return True
|
|
except Exception as e:
|
|
print(f"加载模型类别信息失败: {e}")
|
|
return False
|
|
|
|
@classmethod
|
|
def get_class_name_cn(cls, class_id):
|
|
"""获取类别的中文名称"""
|
|
# 优先使用动态加载的模型类别
|
|
if cls._model_classes is not None:
|
|
original_name = cls._model_classes.get(class_id)
|
|
if original_name:
|
|
# 如果有对应的中文映射则使用，否则使用原始英文名称
|
|
return cls.DEFAULT_CLASS_NAMES_CN.get(class_id, original_name)
|
|
return f"未知类别_{class_id}"
|
|
# 回退到默认中文类别
|
|
return cls.DEFAULT_CLASS_NAMES_CN.get(class_id, f"未知类别_{class_id}")
|
|
|
|
@classmethod
|
|
def get_class_color(cls, class_id):
|
|
"""获取类别对应的颜色 - 始终返回绿色"""
|
|
if cls._model_colors is not None:
|
|
return cls._model_colors.get(class_id, (0, 255, 0))
|
|
return cls.DEFAULT_CLASS_COLORS.get(class_id, (0, 255, 0))
|
|
|
|
@classmethod
|
|
def get_all_classes(cls):
|
|
"""获取所有类别信息"""
|
|
if cls._model_classes is not None:
|
|
return cls._model_classes
|
|
return cls.DEFAULT_CLASS_NAMES_CN
|
|
|
|
# 全局配置实例
|
|
config = DetectionConfig() |