# gpu_parallel.py # -*- coding: utf-8 -*- """ Day 20: GPU 并行推理优化模块 使用 CUDA Stream 让盲道检测和障碍物检测并行执行 """ import os import time import torch import numpy as np from typing import Tuple, Optional, List, Any from concurrent.futures import ThreadPoolExecutor import logging logger = logging.getLogger(__name__) # 全局 CUDA Stream（延迟初始化） _cuda_streams = None _parallel_executor = None def _init_cuda_streams(): """初始化 CUDA Streams""" global _cuda_streams if _cuda_streams is None and torch.cuda.is_available(): try: _cuda_streams = [torch.cuda.Stream() for _ in range(2)] logger.info("[GPU_PARALLEL] 已创建 2 个 CUDA Stream") except Exception as e: logger.warning(f"[GPU_PARALLEL] 创建 CUDA Stream 失败: {e}") _cuda_streams = [] return _cuda_streams def _init_parallel_executor(): """初始化并行执行器（用于 CPU 后处理）""" global _parallel_executor if _parallel_executor is None: _parallel_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="gpu_post") logger.info("[GPU_PARALLEL] 已创建 GPU 后处理线程池") return _parallel_executor class ParallelDetector: """ 并行检测器 - 同时执行盲道检测和障碍物检测 使用方式: detector = ParallelDetector(yolo_model, obstacle_detector) blind_mask, cross_mask, obstacles = detector.detect_all(image, path_mask) """ def __init__(self, yolo_model, obstacle_detector): self.yolo_model = yolo_model self.obstacle_detector = obstacle_detector # 检测参数（从环境变量读取） self.imgsz = int(os.getenv("AIGLASS_YOLO_IMGSZ", "480")) self.use_half = os.getenv("AIGLASS_YOLO_HALF", "1") == "1" self.blind_conf_threshold = 0.20 self.cross_conf_threshold = 0.30 # 初始化 _init_cuda_streams() _init_parallel_executor() logger.info(f"[GPU_PARALLEL] ParallelDetector 初始化完成: imgsz={self.imgsz}, half={self.use_half}") def detect_all( self, image: np.ndarray, path_mask: Optional[np.ndarray] = None ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List[Any]]: """ 并行执行所有检测 Args: image: BGR 图像 path_mask: 盲道掩码（用于障碍物过滤） Returns: (blind_path_mask, crosswalk_mask, obstacles) """ t0 = time.perf_counter() streams = _init_cuda_streams() if streams and len(streams) >= 2: return self._detect_with_streams(image, path_mask, streams) else: # 回退到串行执行 return self._detect_serial(image, path_mask) def _detect_with_streams( self, image: np.ndarray, path_mask: Optional[np.ndarray], streams: List[torch.cuda.Stream] ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List[Any]]: """Day 21: 使用 ThreadPoolExecutor 真正并行检测（替代无效的 CUDA Stream）""" blind_mask = None cross_mask = None obstacles = [] executor = _init_parallel_executor() # 定义两个检测任务 def task_blind_path(): if self.yolo_model is None: return None, None try: results = self.yolo_model.predict( image, verbose=False, conf=min(self.blind_conf_threshold, self.cross_conf_threshold), classes=[0, 1], # 0=crosswalk, 1=blind_path imgsz=self.imgsz, half=self.use_half ) if results and results[0] and results[0].masks is not None: return self._parse_seg_results(results[0], image.shape) except Exception as e: logger.error(f"[GPU_PARALLEL] 盲道检测失败: {e}") return None, None def task_obstacles(): if self.obstacle_detector is None: return [] try: return self.obstacle_detector.detect(image, path_mask=path_mask) except Exception as e: logger.error(f"[GPU_PARALLEL] 障碍物检测失败: {e}") return [] # 并行提交两个任务 from concurrent.futures import as_completed futures = { executor.submit(task_blind_path): 'blind', executor.submit(task_obstacles): 'obstacle' } # 等待所有任务完成 for future in as_completed(futures, timeout=2.0): task_type = futures[future] try: result = future.result() if task_type == 'blind': blind_mask, cross_mask = result else: obstacles = result except Exception as e: logger.error(f"[GPU_PARALLEL] {task_type}任务异常: {e}") return blind_mask, cross_mask, obstacles def _detect_serial( self, image: np.ndarray, path_mask: Optional[np.ndarray] ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List[Any]]: """串行检测（回退模式）""" blind_mask = None cross_mask = None obstacles = [] # 盲道检测 if self.yolo_model is not None: try: results = self.yolo_model.predict( image, verbose=False, conf=min(self.blind_conf_threshold, self.cross_conf_threshold), classes=[0, 1], imgsz=self.imgsz, half=self.use_half ) if results and results[0] and results[0].masks is not None: blind_mask, cross_mask = self._parse_seg_results(results[0], image.shape) except Exception as e: logger.error(f"[GPU_PARALLEL] 盲道检测失败: {e}") # 障碍物检测 if self.obstacle_detector is not None: try: obstacles = self.obstacle_detector.detect(image, path_mask=path_mask) except Exception as e: logger.error(f"[GPU_PARALLEL] 障碍物检测失败: {e}") return blind_mask, cross_mask, obstacles def _parse_seg_results( self, result, image_shape: Tuple[int, int, int] ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: """解析 YOLO 分割结果""" blind_mask = None cross_mask = None h, w = image_shape[:2] if result.masks is None or result.boxes is None: return None, None for mask_tensor, conf_tensor, cls_tensor in zip( result.masks.data, result.boxes.conf, result.boxes.cls ): class_id = int(cls_tensor.item()) confidence = float(conf_tensor.item()) # 置信度过滤 if class_id == 1 and confidence < self.blind_conf_threshold: continue if class_id == 0 and confidence < self.cross_conf_threshold: continue # 转换掩码 current_mask = self._tensor_to_mask(mask_tensor, w, h) if class_id == 1: # 盲道 if blind_mask is None: blind_mask = current_mask else: blind_mask = np.bitwise_or(blind_mask, current_mask) elif class_id == 0: # 斑马线 if cross_mask is None: cross_mask = current_mask else: cross_mask = np.bitwise_or(cross_mask, current_mask) return blind_mask, cross_mask def _tensor_to_mask( self, mask_tensor: torch.Tensor, out_w: int, out_h: int ) -> np.ndarray: """将 PyTorch 张量掩码转换为 NumPy 数组""" import cv2 # 转换为 numpy if mask_tensor.is_cuda: mask_np = mask_tensor.cpu().numpy() else: mask_np = mask_tensor.numpy() # 调整大小 if mask_np.shape[0] != out_h or mask_np.shape[1] != out_w: mask_np = cv2.resize(mask_np, (out_w, out_h), interpolation=cv2.INTER_NEAREST) # 二值化 mask_np = (mask_np > 0.5).astype(np.uint8) * 255 return mask_np def detect_all_parallel( yolo_model, obstacle_detector, image: np.ndarray, path_mask: Optional[np.ndarray] = None ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], List[Any]]: """ 便捷函数：并行执行所有检测 用于替换 workflow_blindpath.py 中的串行检测 """ detector = ParallelDetector(yolo_model, obstacle_detector) return detector.detect_all(image, path_mask)