import logging import traceback from typing import Dict, List, Any, Optional logger = logging.getLogger(__name__) class ObjectExtractor: """ 專門處理物件檢測結果的提取和預處理 負責從YOLO檢測結果提取物件資訊、物件分類和核心物件的辨識 """ def __init__(self, class_names: Dict[int, str] = None, object_categories: Dict[str, List[int]] = None): """ 初始化物件提取器 Args: class_names: 類別ID到類別名稱的映射字典 object_categories: 物件類別分組字典 """ try: self.class_names = class_names or {} self.object_categories = object_categories or {} # 1. 讀取並設定基本信心度門檻（如果外部沒傳，就預設 0.25） self.base_conf_threshold = 0.25 # 2. 動態信心度調整映射表 (key: 小寫 class_name, value: 調整係數) # 最終的門檻 = base_conf_threshold * factor # 如果某個 class_name 沒在這裡，就直接用 base_conf_threshold（相當於 factor=1.0） self.dynamic_conf_map = { "traffic light": 0.6, "car": 0.8, "person": 0.7, } logger.info(f"ObjectExtractor initialized with {len(self.class_names)} class names and {len(self.object_categories)} object categories") except Exception as e: logger.error(f"Failed to initialize ObjectExtractor: {str(e)}") logger.error(traceback.format_exc()) raise def _get_dynamic_threshold(self, class_name: str) -> float: """ 根據 class_name 從 dynamic_conf_map 拿到 factor，計算最終的信心度門檻： threshold = base_conf_threshold * factor 如果 class_name 不在映射表裡，就回傳 base_conf_threshold。 """ # 使用小寫做匹配，確保在 dynamic_conf_map 裡的 key 也都用小寫 key = class_name.lower() factor = self.dynamic_conf_map.get(key, 1.0) return self.base_conf_threshold * factor def extract_detected_objects( self, detection_result: Any, confidence_threshold: float = 0.25, region_analyzer=None ) -> List[Dict]: """ 從檢測結果中提取物件資訊，包含位置資訊 Args: detection_result: YOLO檢測結果 confidence_threshold: 改由動態門檻決定 region_analyzer: 區域分析器實例，用於判斷物件所屬區域 Returns: 包含檢測物件資訊的字典列表 """ try: # 調試信息：記錄當前類別映射狀態 logger.info(f"ObjectExtractor.extract_detected_objects called") logger.info(f"Current class_names keys: {list(self.class_names.keys()) if self.class_names else 'None'}") if detection_result is None: logger.warning("Detection result is None") return [] if not hasattr(detection_result, 'boxes'): logger.error("Detection result does not have boxes attribute") return [] boxes = detection_result.boxes.xyxy.cpu().numpy() classes = detection_result.boxes.cls.cpu().numpy().astype(int) confidences = detection_result.boxes.conf.cpu().numpy() # 獲取圖像尺寸 img_height, img_width = detection_result.orig_shape[:2] detected_objects = [] for box, class_id, confidence in zip(boxes, classes, confidences): try: # 1. 先拿到這筆偵測物件的 class_name class_name = self.class_names.get(int(class_id), f"unknown_class_{class_id}") # 2. 計算這個 class 應該採用的動態 threshold dyn_thr = self._get_dynamic_threshold(class_name) # e.g. 0.25 * factor # 3. 如果 confidence < dyn_thr，就跳過這一筆 if confidence < dyn_thr: continue # 後面維持原本的座標、中心、大小、區域等資訊計算 x1, y1, x2, y2 = box width = x2 - x1 height = y2 - y1 # 中心點計算 center_x = (x1 + x2) / 2 center_y = (y1 + y2) / 2 # 標準化位置 (0-1) norm_x = center_x / img_width norm_y = center_y / img_height norm_width = width / img_width norm_height = height / img_height # 面積計算 area = width * height norm_area = area / (img_width * img_height) # 區域判斷 object_region = "unknown" if region_analyzer: object_region = region_analyzer.determine_region(norm_x, norm_y) # 調試信息：記錄映射過程 if class_name.startswith("unknown_class_"): logger.warning( f"Class ID {class_id} not found in class_names. " f"Available keys: {list(self.class_names.keys())}" ) else: logger.debug(f"Successfully mapped class ID {class_id} to '{class_name}'") detected_objects.append({ "class_id": int(class_id), "class_name": class_name, "confidence": float(confidence), "box": [float(x1), float(y1), float(x2), float(y2)], "center": [float(center_x), float(center_y)], "normalized_center": [float(norm_x), float(norm_y)], "size": [float(width), float(height)], "normalized_size": [float(norm_width), float(norm_height)], "area": float(area), "normalized_area": float(norm_area), "region": object_region }) except Exception as e: logger.error(f"Error processing object with class_id {class_id}: {str(e)}") continue logger.info(f"Extracted {len(detected_objects)} objects from detection result") # print(f"DEBUG: ObjectExtractor filtered objects by class:") # for class_name in ["car", "traffic light", "person", "handbag"]: # class_objects = [obj for obj in detected_objects if obj.get("class_name") == class_name] # if class_objects: # confidences = [obj.get("confidence", 0) for obj in class_objects] # print(f"DEBUG: {class_name}: {len(class_objects)} objects, confidences: {confidences}") # print(f"DEBUG: base_conf_threshold: {self.base_conf_threshold}") # print(f"DEBUG: dynamic_conf_map: {self.dynamic_conf_map}") return detected_objects except Exception as e: logger.error(f"Error extracting detected objects: {str(e)}") logger.error(traceback.format_exc()) return [] def update_class_names(self, class_names: Dict[int, str]): """ 動態更新類別名稱映射 Args: class_names: 新的類別名稱映射字典 """ try: self.class_names = class_names or {} logger.info(f"Class names updated: {len(self.class_names)} classes") logger.debug(f"Updated class names: {self.class_names}") except Exception as e: logger.error(f"Failed to update class names: {str(e)}") def categorize_object(self, obj: Dict) -> str: """ 將檢測到的物件分類到功能類別中，用於區域識別 Args: obj: 物件字典 Returns: 物件功能類別字串 """ try: class_id = obj.get("class_id", -1) class_name = obj.get("class_name", "").lower() # 使用現有的類別映射（如果可用） if self.object_categories: for category, ids in self.object_categories.items(): if class_id in ids: return category # 基於COCO類別名稱的後備分類 furniture_items = ["chair", "couch", "bed", "dining table", "toilet"] plant_items = ["potted plant"] electronic_items = ["tv", "laptop", "mouse", "remote", "keyboard", "cell phone"] vehicle_items = ["bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat"] person_items = ["person"] kitchen_items = ["bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "refrigerator", "oven", "toaster", "sink", "microwave"] sports_items = ["frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket"] personal_items = ["handbag", "tie", "suitcase", "umbrella", "backpack"] if any(item in class_name for item in furniture_items): return "furniture" elif any(item in class_name for item in plant_items): return "plant" elif any(item in class_name for item in electronic_items): return "electronics" elif any(item in class_name for item in vehicle_items): return "vehicle" elif any(item in class_name for item in person_items): return "person" elif any(item in class_name for item in kitchen_items): return "kitchen_items" elif any(item in class_name for item in sports_items): return "sports" elif any(item in class_name for item in personal_items): return "personal_items" else: return "misc" except Exception as e: logger.error(f"Error categorizing object: {str(e)}") logger.error(traceback.format_exc()) return "misc" def get_object_categories(self, detected_objects: List[Dict]) -> set: """ 從檢測到的物件中取得唯一的物件類別 Args: detected_objects: 檢測到的物件列表 Returns: 唯一物件類別的集合 """ try: object_categories = set() for obj in detected_objects: category = self.categorize_object(obj) if category: object_categories.add(category) logger.info(f"Found {len(object_categories)} unique object categories") return object_categories except Exception as e: logger.error(f"Error getting object categories: {str(e)}") logger.error(traceback.format_exc()) return set() def identify_core_objects_for_scene(self, detected_objects: List[Dict], scene_type: str) -> List[Dict]: """ 識別定義特定場景類型的核心物件 Args: detected_objects: 檢測到的物件列表 scene_type: 場景類型 Returns: 場景的核心物件列表 """ try: core_objects = [] # 場景核心物件映射 scene_core_mapping = { "bedroom": [59], # bed "kitchen": [68, 69, 71, 72], # microwave, oven, sink, refrigerator "living_room": [57, 58, 62], # sofa, chair, tv "dining_area": [60, 42, 43], # dining table, fork, knife "office_workspace": [63, 64, 66, 73] # laptop, mouse, keyboard, book } if scene_type in scene_core_mapping: core_class_ids = scene_core_mapping[scene_type] for obj in detected_objects: if obj.get("class_id") in core_class_ids and obj.get("confidence", 0) >= 0.4: core_objects.append(obj) logger.info(f"Identified {len(core_objects)} core objects for scene type '{scene_type}'") return core_objects except Exception as e: logger.error(f"Error identifying core objects for scene '{scene_type}': {str(e)}") logger.error(traceback.format_exc()) return [] def group_objects_by_category_and_region(self, detected_objects: List[Dict]) -> Dict: """ 將物件按類別和區域分組 Args: detected_objects: 檢測到的物件列表 Returns: 按類別和區域分組的物件字典 """ try: category_regions = {} for obj in detected_objects: category = self.categorize_object(obj) if not category: continue if category not in category_regions: category_regions[category] = {} region = obj.get("region", "center") if region not in category_regions[category]: category_regions[category][region] = [] category_regions[category][region].append(obj) logger.info(f"Grouped objects into {len(category_regions)} categories across regions") return category_regions except Exception as e: logger.error(f"Error grouping objects by category and region: {str(e)}") logger.error(traceback.format_exc()) return {} def filter_objects_by_confidence(self, detected_objects: List[Dict], min_confidence: float) -> List[Dict]: """ 根據信心度過濾物件 Args: detected_objects: 檢測到的物件列表 min_confidence: 最小信心度閾值 Returns: 過濾後的物件列表 """ try: filtered_objects = [ obj for obj in detected_objects if obj.get("confidence", 0) >= min_confidence ] logger.info(f"Filtered {len(detected_objects)} objects to {len(filtered_objects)} objects with confidence >= {min_confidence}") return filtered_objects except Exception as e: logger.error(f"Error filtering objects by confidence: {str(e)}") logger.error(traceback.format_exc()) return detected_objects # 發生錯誤時返回原始列表