Pārlūkot izejas kodu

修复隐患识别的拼音返回问题
新增次二级场景对应二级再对应三级的功能

KCY 1 mēnesi atpakaļ
vecāks
revīzija
1cbaae2695

+ 35 - 1
shudao-chat-py/routers/hazard.py

@@ -20,6 +20,7 @@ from models.scene import RecognitionRecord, Scene, SecondScene, ThirdScene
 from services.oss_service import oss_service
 from services.yolo_service import yolo_service
 from utils.crypto import decrypt_url
+from utils.label_translator import translate_labels, translate_scenes_for_query
 from utils.logger import logger
 
 router = APIRouter()
@@ -207,22 +208,43 @@ def _unique_ordered(items: List[str]) -> List[str]:
 def _build_scene_relations(
     db: Session, normalized_labels: List[str]
 ) -> tuple[list[str], dict[str, list[str]]]:
+    """
+    Build scene relations by mapping sub-secondary scenes to secondary scenes,
+    then querying for third scenes (hazards).
+    
+    Args:
+        db: Database session
+        normalized_labels: List of sub-secondary scene names (e.g., "加油机", "加油枪")
+    
+    Returns:
+        Tuple of (third_scene_names, element_hazards)
+        - third_scene_names: List of all third scene names (hazards)
+        - element_hazards: Dict mapping original labels to their hazards
+    """
     third_scene_names: List[str] = []
     element_hazards: Dict[str, List[str]] = {}
 
     for label in normalized_labels:
+        # Translate sub-secondary scene to secondary scene for database query
+        # e.g., "加油机" -> "加油设施及附属场地"
+        query_label = translate_scenes_for_query([label])[0] if label else label
+        
+        # Query SecondScene table using the translated secondary scene name
         second_scene = (
             db.query(SecondScene)
             .filter(
-                SecondScene.second_scene_name == label,
+                SecondScene.second_scene_name == query_label,
                 SecondScene.is_deleted == 0,
             )
             .first()
         )
+        
         if not second_scene:
+            # No matching secondary scene found, store empty hazards list
             element_hazards.setdefault(label, [])
             continue
 
+        # Query ThirdScene table for hazards related to this secondary scene
         third_scenes = (
             db.query(ThirdScene)
             .filter(
@@ -231,9 +253,11 @@ def _build_scene_relations(
             )
             .all()
         )
+        
         if third_scenes:
             current_element_hazards = [item.third_scene_name for item in third_scenes]
             third_scene_names.extend(current_element_hazards)
+            # Store hazards under the ORIGINAL sub-secondary scene label
             element_hazards[label] = _dedupe_list(
                 element_hazards.get(label, []) + current_element_hazards
             )
@@ -424,6 +448,16 @@ async def hazard(
             labels,
             normalized_labels,
         )
+        
+        # Translate pinyin/English labels to Chinese
+        translated_labels = translate_labels(normalized_labels, fallback_to_original=True)
+        logger.info(
+            "[hazard][%s] translated labels: before=%s, after=%s",
+            request_id,
+            normalized_labels,
+            translated_labels,
+        )
+        normalized_labels = translated_labels
 
         detection_results = []
         for i, label in enumerate(normalized_labels):

+ 268 - 0
shudao-chat-py/utils/label_translator.py

@@ -0,0 +1,268 @@
+"""
+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",
+]