import os import cv2 import time import logging import threading import numpy as np class VisionDetector: """ 通用视觉检测设施 内建两把锁，确保多线程下模型加载与GPU推理绝对安全 """ def __init__(self, image_getter=None): self.image_getter = image_getter self.yolo_models = {} # 加载锁：防止多线程同时挤破头去加载同一个模型把显存撑爆 self._model_load_lock = threading.Lock() # 推理锁：保证当前只有一个人在用GPU算图，避免多线程抢占卡死 self._inference_lock = threading.Lock() def load_model(self, model_path: str) -> bool: """ 安全地把模型揣进内存里，已经有了就不重复拿。 """ if not model_path or not os.path.exists(model_path): logging.error(f"找不到模型文件: {model_path}") return False with self._model_load_lock: if model_path not in self.yolo_models: try: from ultralytics import YOLO logging.info(f"正在加载模型: {model_path}") self.yolo_models[model_path] = YOLO(model_path) return True except Exception as e: logging.error(f"加载模型失败: {str(e)}") return False return True def get_camera_image(self, timeout=0.5): """从硬件拿一张新鲜的画面。""" import rclpy if not rclpy.ok(): return None if not self.image_getter: return None try: image, _ = self.image_getter.get_image(timeout=timeout) return image except Exception as e: logging.error(f"拿图时出了点岔子: {e}") return None def detect_color_hsv(self, image: np.ndarray, roi_coords: list) -> tuple: """给指定框里的颜色做个体检，看看大面上是什么颜色。""" if image is None or not roi_coords: return "未知颜色", 0, 0, 0 img_h, img_w = image.shape[:2] x, y, w, h = roi_coords x = max(0, min(x, img_w - 1)) y = max(0, min(y, img_h - 1)) w = max(1, min(w, img_w - x)) h = max(1, min(h, img_h - y)) roi = image[y:y+h, x:x+w] if roi.size == 0: return "无效范围", 0, 0, 0 avg_color = np.mean(roi, axis=(0, 1)) avg_b, avg_g, avg_r = avg_color rgb_color = np.uint8([[[avg_b, avg_g, avg_r]]]) hsv_color = cv2.cvtColor(rgb_color, cv2.COLOR_BGR2HSV) h_val, s_val, v_val = hsv_color[0][0] color_ranges = { "红色": [(0, 10), (170, 180)], "橙色": [(11, 25)], "黄色": [(26, 35)], "绿色": [(36, 85)], "青色": [(86, 100)], "蓝色": [(101, 130)], "紫色": [(131, 150)], "粉色": [(151, 169)], } if v_val < 30: color_name = "黑色" elif s_val < 40: if v_val > 200: color_name = "白色" elif v_val > 150: color_name = "亮灰色" elif v_val > 80: color_name = "灰色" else: color_name = "深灰色" elif v_val < 60: color_name = "暗色" else: color_name = "未知颜色" for c_name, ranges in color_ranges.items(): for h_range in ranges: if h_range[0] <= h_val <= h_range[1]: if c_name == "紫色" and s_val < 80: color_name = "灰色偏紫" else: color_name = c_name break return color_name, int(avg_r), int(avg_g), int(avg_b) def run_yolo_inference(self, image: np.ndarray, model_path: str) -> list: """全图跑一遍推理。这是算力消耗大户，已上锁防冲突。""" if model_path not in self.yolo_models: if not self.load_model(model_path): return [] model = self.yolo_models[model_path] detections = [] try: with self._inference_lock: results = model(image, verbose=False) for result in results: boxes = result.boxes if boxes is None or len(boxes) == 0: continue boxes_data = boxes.xyxy.cpu().numpy() classes_data = boxes.cls.cpu().numpy() confidences_data = boxes.conf.cpu().numpy() for i in range(len(boxes_data)): x1, y1, x2, y2 = boxes_data[i] detections.append({ 'class': result.names[int(classes_data[i])], 'confidence': float(confidences_data[i]), 'bbox': [int(x1), int(y1), int(x2), int(y2)], 'center': [int((x1+x2)/2), int((y1+y2)/2)] }) except Exception as e: logging.error(f"推理图的时候卡壳了: {str(e)}") return detections def filter_objects_by_static_roi(self, objects: list, roi_coords: list, target_class: str = None, conf_thresh: float = 0.5) -> list: """筛查出中心点落在指定死框里，并且符合身份的目标。""" rx, ry, rw, rh = roi_coords valid = [] for obj in objects: if obj['confidence'] < conf_thresh: continue if target_class and target_class not in obj['class']: continue cx, cy = obj['center'] if (rx <= cx <= rx + rw) and (ry <= cy <= ry + rh): valid.append(obj) return valid def generate_dynamic_rois(self, objects: list, ref_class: str, offset_dx: int, offset_dy: int, roi_w: int, roi_h: int, img_w: int, img_h: int) -> list: """以某个参考物为锚点，往旁边偏移一段距离，圈出几个活框来。""" dynamic_rois = [] for obj in objects: if ref_class in obj['class']: cx, cy = obj['center'] rx = max(0, min(int(cx + offset_dx - roi_w / 2), img_w - roi_w)) ry = max(0, min(int(cy + offset_dy - roi_h / 2), img_h - roi_h)) dynamic_rois.append([rx, ry, roi_w, roi_h]) return dynamic_rois def draw_detections(self, image: np.ndarray, detections: list, roi_coords: list = None, is_dynamic: bool = False, ref_rois: list = None) -> np.ndarray: """在图像上画出检测目标和ROI区域，支持同时画多个ROI框以满足并行检测需求。""" if image is None: return None img_draw = image.copy() # 1. 画死框 (红色) if roi_coords and not is_dynamic: # 判断传入的是单个框 [x,y,w,h] 还是多个框 [[x,y,w,h], ...] if len(roi_coords) > 0 and isinstance(roi_coords[0], (int, float)): rois_to_draw = [roi_coords] else: rois_to_draw = roi_coords for rx, ry, rw, rh in rois_to_draw: cv2.rectangle(img_draw, (rx, ry), (rx + rw, ry + rh), (0, 0, 255), 2) cv2.putText(img_draw, "ROI", (rx, ry - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) # 2. 画活框的区域 (蓝色) if is_dynamic and ref_rois: for r in ref_rois: rx, ry, rw, rh = r cv2.rectangle(img_draw, (rx, ry), (rx + rw, ry + rh), (255, 0, 0), 2) cv2.putText(img_draw, "Dyn_ROI", (rx, ry - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2) # 3. 画目标框 (绿色) for obj in detections: x1, y1, x2, y2 = obj['bbox'] cv2.rectangle(img_draw, (x1, y1), (x2, y2), (0, 255, 0), 2) label = f"{obj['class']} {obj['confidence']:.2f}" cv2.putText(img_draw, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) return img_draw def save_image_to_disk(self, image: np.ndarray, save_dir: str, prefix: str, suffix: str, max_files: int = 200) -> bool: """把图存硬盘里留底，并限制文件夹内图片最大数量。""" try: os.makedirs(save_dir, exist_ok=True) timestamp = time.strftime("%Y%m%d_%H%M%S") filename = f"{timestamp}_{prefix}_{suffix}.jpg" filepath = os.path.join(save_dir, filename) cv2.imwrite(filepath, image) # 限制最大保存数量（200张限制机制） if max_files > 0: files = [os.path.join(save_dir, f) for f in os.listdir(save_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))] if len(files) > max_files: # 按文件修改时间排序，最旧的文件排在最前面 files.sort(key=os.path.getmtime) for f in files[:-max_files]: try: os.remove(f) except Exception: pass return True except Exception as e: logging.error(f"图存不下来: {str(e)}") return False