shudao-main/shudao-chat-py/utils/label_translator.py

"""
Label translator for hazard detection system.
Translates pinyin/English labels to Chinese for display and database queries.
"""

from typing import List, Dict, Optional


# Complete mapping: pinyin/English -> Chinese
LABEL_MAPPING: Dict[str, str] = {
    # Gas Station (加油站) labels
    "fangzhuanglan": "防撞栏",
    "jiayouqiang": "加油枪",
    "jiayouji": "加油机",
    "miehuoqi": "灭火器",
    "zhuixingtong": "锥形桶",
    "push_extinguisher": "手推灭火器",
    "support_pillar": "支撑柱",
    "fuel_hose": "加油机输油软管",
    "vent_pipe": "通气管",
    "oil_discharge_port": "卸油口",
    "static_release_post": "静电释放桩",
    "static_wire": "静电线",
    "static_clip": "静电夹",
    "power_distribution_box": "配电箱",
    
    # Alternative English names (aliases)
    "crash_barrier": "防撞栏",
    "fuel_nozzle": "加油枪",
    "fuel_dispenser": "加油机",
    "fire_extinguisher": "灭火器",
    "traffic_cone": "锥形桶",
    "pillar": "支撑柱",
    "oil_hose": "加油机输油软管",
    "ventilation_pipe": "通气管",
    "discharge_port": "卸油口",
    "grounding_post": "静电释放桩",
    "grounding_wire": "静电线",
    "grounding_clip": "静电夹",
    "distribution_box": "配电箱",
}

# Reverse mapping: Chinese -> pinyin/English (for reference)
REVERSE_MAPPING: Dict[str, str] = {v: k for k, v in LABEL_MAPPING.items()}


def translate_label(label: str, fallback_to_original: bool = True) -> str:
    """
    Translate a single label from pinyin/English to Chinese.
    
    Args:
        label: The label to translate (pinyin or English)
        fallback_to_original: If True, return original label when no translation found
        
    Returns:
        Translated Chinese label, or original label if not found and fallback is True
        
    Examples:
        >>> translate_label("fangzhuanglan")
        "防撞栏"
        >>> translate_label("fuel_hose")
        "加油机输油软管"
        >>> translate_label("防撞栏")  # Already Chinese
        "防撞栏"
        >>> translate_label("unknown_label")
        "unknown_label"
    """
    if not label:
        return label
    
    label_stripped = label.strip()
    if not label_stripped:
        return label
    
    # Check if already Chinese (contains Chinese characters)
    if _contains_chinese(label_stripped):
        return label_stripped
    
    # Try exact match (case-insensitive)
    label_lower = label_stripped.lower()
    if label_lower in LABEL_MAPPING:
        return LABEL_MAPPING[label_lower]
    
    # Try with underscores replaced by spaces
    label_spaced = label_lower.replace("_", " ")
    if label_spaced in LABEL_MAPPING:
        return LABEL_MAPPING[label_spaced]
    
    # Fallback to original if no translation found
    if fallback_to_original:
        return label_stripped
    
    return ""


def translate_labels(labels: List[str], fallback_to_original: bool = True) -> List[str]:
    """
    Translate a list of labels from pinyin/English to Chinese.
    
    Args:
        labels: List of labels to translate
        fallback_to_original: If True, keep original labels when no translation found
        
    Returns:
        List of translated Chinese labels
        
    Examples:
        >>> translate_labels(["fangzhuanglan", "jiayouqiang", "防撞栏"])
        ["防撞栏", "加油枪", "防撞栏"]
    """
    if not labels:
        return []
    
    return [translate_label(label, fallback_to_original) for label in labels]


def is_translated(label: str) -> bool:
    """
    Check if a label is already in Chinese (translated).
    
    Args:
        label: The label to check
        
    Returns:
        True if label contains Chinese characters, False otherwise
    """
    return _contains_chinese(label)


def get_english_name(chinese_label: str) -> Optional[str]:
    """
    Get the English/pinyin name for a Chinese label (reverse lookup).
    
    Args:
        chinese_label: Chinese label to look up
        
    Returns:
        English/pinyin name if found, None otherwise
        
    Examples:
        >>> get_english_name("防撞栏")
        "fangzhuanglan"
    """
    return REVERSE_MAPPING.get(chinese_label.strip())


def add_custom_mapping(pinyin_or_english: str, chinese: str) -> None:
    """
    Add a custom label mapping at runtime.
    
    Args:
        pinyin_or_english: The pinyin or English label
        chinese: The Chinese translation
        
    Note:
        This is useful for dynamic label additions without modifying the source code.
    """
    key = pinyin_or_english.strip().lower()
    value = chinese.strip()
    LABEL_MAPPING[key] = value
    REVERSE_MAPPING[value] = key


def get_all_mappings() -> Dict[str, str]:
    """
    Get all label mappings.
    
    Returns:
        Dictionary of all pinyin/English -> Chinese mappings
    """
    return LABEL_MAPPING.copy()


def _contains_chinese(text: str) -> bool:
    """
    Check if text contains Chinese characters.
    
    Args:
        text: Text to check
        
    Returns:
        True if text contains at least one Chinese character
    """
    if not text:
        return False
    
    for char in text:
        if '\u4e00' <= char <= '\u9fff':
            return True
    
    return False


# Sub-secondary to secondary scene mapping
# Used only for database queries - these scenes still appear as-is in UI and storage
SUB_SECONDARY_TO_SECONDARY_MAPPING: Dict[str, str] = {
    # 加油设施_附属场地 -> maps to parent secondary scene (note: uses underscore in DB)
    "加油机": "加油设施_附属场地",
    "加油枪": "加油设施_附属场地",
    "加油机输油软管": "加油设施_附属场地",
    "防撞栏": "加油设施_附属场地",
    "锥形桶": "加油设施_附属场地",
    "灭火器": "加油设施_附属场地",
    
    # 罩棚 and 立柱 are separate scenes in DB, map support pillar to 立柱
    "支撑立柱": "立柱",
}


def translate_scene_for_query(scene_name: str) -> str:
    """
    Translate sub-secondary scene to its parent secondary scene for database queries.
    This is ONLY used when querying the database, not for display or storage.
    
    Args:
        scene_name: The scene name to translate (could be sub-secondary or secondary)
        
    Returns:
        Parent secondary scene name if it's a sub-secondary scene, otherwise original name
        
    Examples:
        >>> translate_scene_for_query("加油机")
        "加油设施及附属场地"
        >>> translate_scene_for_query("支撑立柱")
        "罩棚及立柱"
        >>> translate_scene_for_query("加油设施及附属场地")
        "加油设施及附属场地"
    """
    if not scene_name:
        return scene_name
    
    scene_stripped = scene_name.strip()
    return SUB_SECONDARY_TO_SECONDARY_MAPPING.get(scene_stripped, scene_stripped)


def translate_scenes_for_query(scene_names: List[str]) -> List[str]:
    """
    Translate a list of sub-secondary scenes to their parent secondary scenes for database queries.
    
    Args:
        scene_names: List of scene names to translate
        
    Returns:
        List of translated scene names (parent secondary scenes)
        
    Examples:
        >>> translate_scenes_for_query(["加油机", "加油枪", "支撑立柱"])
        ["加油设施及附属场地", "加油设施及附属场地", "罩棚及立柱"]
    """
    if not scene_names:
        return []
    
    return [translate_scene_for_query(scene) for scene in scene_names]


# Export main functions
__all__ = [
    "translate_label",
    "translate_labels",
    "is_translated",
    "get_english_name",
    "add_custom_mapping",
    "get_all_mappings",
    "translate_scene_for_query",
    "translate_scenes_for_query",
    "LABEL_MAPPING",
    "SUB_SECONDARY_TO_SECONDARY_MAPPING",
]