wrz-yolo/config.py

#!/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: "白线模糊"
    }
    
    # 通用类别关键词映射（扩展版）
    COMMON_CLASS_MAPPINGS = {
        # 道路损伤相关
        'crack': '裂缝',
        'pothole': '坑洞', 
        'damage': '损伤',
        'longitudinal': '纵向裂缝',
        'transverse': '横向裂缝',
        'alligator': '网状裂缝',
        'rutting': '车辙',
        'bleeding': '泛油',
        'weathering': '风化',
        'patching': '修补',
        
        # 交通工具
        'person': '人',
        'people': '人群',
        'car': '汽车',
        'truck': '卡车',
        'bus': '公交车',
        'motorcycle': '摩托车',
        'bicycle': '自行车',
        'bike': '自行车',
        'vehicle': '车辆',
        
        # 交通设施
        'road': '道路',
        'traffic': '交通',
        'sign': '标志',
        'light': '信号灯',
        'signal': '信号',
        'stop': '停车标志',
        'yield': '让行标志',
        'crosswalk': '人行横道',
        'lane': '车道',
        'marking': '标线',
        'white': '白线',
        'yellow': '黄线',
        
        # 其他常见目标
        'animal': '动物',
        'dog': '狗',
        'cat': '猫',
        'bird': '鸟',
        'tree': '树',
        'building': '建筑',
        'pole': '杆子',
        'fence': '围栏'
    }
    
    # 动态颜色生成配置
    COLOR_PALETTE = [
        (0, 255, 0),    # 绿色
        (255, 0, 0),    # 红色  
        (0, 0, 255),    # 蓝色
        (255, 255, 0),  # 黄色
        (255, 0, 255),  # 紫色
        (0, 255, 255),  # 青色
        (255, 128, 0),  # 橙色
        (128, 0, 255),  # 紫罗兰
        (255, 192, 203), # 粉色
        (128, 128, 128)  # 灰色
    ]
    
    # 默认类别颜色配置 (BGR格式)
    DEFAULT_CLASS_COLORS = {
        0: (0, 255, 0),    # 绿色
        1: (255, 0, 0),    # 红色
        2: (0, 0, 255),    # 蓝色
        3: (255, 255, 0),  # 黄色
        4: (255, 0, 255)   # 紫色
    }
    
    @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 = {}
            for i in range(len(model.names)):
                # 使用调色板循环分配颜色
                color_index = i % len(cls.COLOR_PALETTE)
                cls._model_colors[i] = cls.COLOR_PALETTE[color_index]
            
            print(f"成功加载 {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 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):
        """获取类别对应的颜色"""
        # 如果class_id是字符串，尝试转换为整数或使用hash
        if isinstance(class_id, str):
            # 尝试从类别名称获取ID
            if cls._model_classes and class_id in cls._model_classes.values():
                # 找到对应的ID
                for id_val, name in cls._model_classes.items():
                    if name == class_id:
                        class_id = id_val
                        break
            else:
                # 使用hash生成一个稳定的索引
                class_id = hash(class_id) % len(cls.COLOR_PALETTE)
        
        # 确保class_id是整数
        if not isinstance(class_id, int):
            class_id = 0
        
        if cls._model_colors is not None:
            return cls._model_colors.get(class_id, cls.COLOR_PALETTE[class_id % len(cls.COLOR_PALETTE)])
        return cls.DEFAULT_CLASS_COLORS.get(class_id, cls.COLOR_PALETTE[class_id % len(cls.COLOR_PALETTE)])
    
    @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()