import pyzed.sl as sl import cv2 import numpy as np import os import time from scipy.spatial import cKDTree # Try to import fish detector (optional) try: from detector import FishDetector DETECTOR_AVAILABLE = True except ImportError: DETECTOR_AVAILABLE = False FishDetector = None class ZEDReader: def __init__(self, svo_path=None, camera_mode=False, use_yolo_detector=False, yolo_model_path=None): # 创建ZED相机对象 self.zed = sl.Camera() self.camera_mode = camera_mode self.use_yolo_detector = use_yolo_detector self.detector = None # Initialize YOLO detector if requested if use_yolo_detector: if not DETECTOR_AVAILABLE: print("警告: YOLO检测器不可用。请安装detector.py和相关依赖") self.use_yolo_detector = False else: try: self.detector = FishDetector(model_path=yolo_model_path) print("✓ YOLO鱼类检测器已初始化") except Exception as e: print(f"警告: 初始化YOLO检测器失败: {e}") self.use_yolo_detector = False # 设置初始化参数 self.init_params = sl.InitParameters() if svo_path and not camera_mode: # 如果提供SVO文件路径，则从文件读取 print(f"从SVO文件读取: {svo_path}") self.init_params.set_from_svo_file(svo_path) self.init_params.svo_real_time_mode = False else: # 实时相机模式设置 print("使用实时相机模式") self.init_params.camera_resolution = sl.RESOLUTION.HD1200 self.init_params.depth_mode = sl.DEPTH_MODE.NEURAL self.init_params.coordinate_units = sl.UNIT.MILLIMETER self.init_params.camera_fps = 30 # 设置帧率 def open(self): # 打开相机或SVO文件 err = self.zed.open(self.init_params) if err != sl.ERROR_CODE.SUCCESS: print(f"错误: {err}") return False # 获取并打印标定参数 self.print_calibration_params() return True def print_calibration_params(self): """打印相机标定参数""" calibration_params = self.zed.get_camera_information().camera_configuration.calibration_parameters print("\n=== 相机标定参数 ===") # 左眼相机参数 print("\n左眼相机参数:") print(f"焦距: fx={calibration_params.left_cam.fx}, fy={calibration_params.left_cam.fy}") print(f"主点: cx={calibration_params.left_cam.cx}, cy={calibration_params.left_cam.cy}") print("畸变参数:") print(f"k1={calibration_params.left_cam.disto[0]}") print(f"k2={calibration_params.left_cam.disto[1]}") print(f"k3={calibration_params.left_cam.disto[2]}") print(f"k4={calibration_params.left_cam.disto[3]}") print(f"k5={calibration_params.left_cam.disto[4]}") # 右眼相机参数 print("\n右眼相机参数:") print(f"焦距: fx={calibration_params.right_cam.fx}, fy={calibration_params.right_cam.fy}") print(f"主点: cx={calibration_params.right_cam.cx}, cy={calibration_params.right_cam.cy}") print("畸变参数:") print(f"k1={calibration_params.right_cam.disto[0]}") print(f"k2={calibration_params.right_cam.disto[1]}") print(f"k3={calibration_params.right_cam.disto[2]}") print(f"k4={calibration_params.right_cam.disto[3]}") print(f"k5={calibration_params.right_cam.disto[4]}") #import pdb; pdb.set_trace() print("\n立体参数:") # 获取旋转矩阵 rotation = calibration_params.stereo_transform.m print("旋转矩阵 R:") print(f"[{rotation[0][0]}, {rotation[0][1]}, {rotation[0][2]}]") print(f"[{rotation[1][0]}, {rotation[1][1]}, {rotation[1][2]}]") print(f"[{rotation[2][0]}, {rotation[2][1]}, {rotation[2][2]}]") print("\n平移向量 T:") print(f"[{rotation[0][3]}, {rotation[1][3]}, {rotation[2][3]}]") # 保存标定参数到文件 if hasattr(self, 'save_path') and self.save_path: self.save_calibration_params(calibration_params) def save_calibration_params(self, calibration_params): """将标定参数保存到文件""" if not hasattr(self, 'save_path') or not self.save_path: return calib_file = os.path.join(self.save_path, 'calibration_params.txt') with open(calib_file, 'w') as f: f.write("=== 相机标定参数 ===\n") # 左眼相机 f.write("\n左眼相机参数:\n") f.write(f"焦距: fx={calibration_params.left_cam.fx}, fy={calibration_params.left_cam.fy}\n") f.write(f"主点: cx={calibration_params.left_cam.cx}, cy={calibration_params.left_cam.cy}\n") f.write("畸变参数:\n") for i, d in enumerate(calibration_params.left_cam.disto): f.write(f"k{i+1}={d}\n") # 右眼相机 f.write("\n右眼相机参数:\n") f.write(f"焦距: fx={calibration_params.right_cam.fx}, fy={calibration_params.right_cam.fy}\n") f.write(f"主点: cx={calibration_params.right_cam.cx}, cy={calibration_params.right_cam.cy}\n") f.write("畸变参数:\n") for i, d in enumerate(calibration_params.right_cam.disto): f.write(f"k{i+1}={d}\n") # 立体参数 f.write("\n立体参数:\n") rotation = calibration_params.stereo_transform.m f.write("旋转矩阵 R:\n") f.write(f"[{rotation[0][0]}, {rotation[0][1]}, {rotation[0][2]}]\n") f.write(f"[{rotation[1][0]}, {rotation[1][1]}, {rotation[1][2]}]\n") f.write(f"[{rotation[2][0]}, {rotation[2][1]}, {rotation[2][2]}]\n") f.write("\n平移向量 T:\n") f.write(f"[{rotation[0][3]}, {rotation[1][3]}, {rotation[2][3]}]\n") def start_recording(self, filename=None): """开始录制""" if filename is None: # 使用时间戳作为文件名 timestamp = time.strftime("%Y%m%d_%H%M%S") filename = f"recording_{timestamp}.svo" recording_param = sl.RecordingParameters(filename, sl.SVO_COMPRESSION_MODE.H264) err = self.zed.enable_recording(recording_param) if err != sl.ERROR_CODE.SUCCESS: print(f"无法开始录制: {err}") return False print(f"开始录制到文件: {filename}") return True def ensure_save_path(self, save_path): """确保保存路径存在，如果不存在则创建""" if save_path: if not os.path.exists(save_path): os.makedirs(save_path) print(f"创建保存目录: {save_path}") return True return False def detect_fish(self, depth_np, calib_params, min_points=100, max_depth=1200, max_y=150): """ 检测深度图中是否有鱼类 检测条件： - z+ 方向（深度）：0-max_depth mm - y- 方向（相机坐标系y轴负方向，y < max_y） - 满足条件的点数 >= min_points Args: depth_np: 深度图的numpy数组 (H, W) 单位为mm calib_params: 相机标定参数 min_points: 检测到鱼所需的最少点数，默认100 max_depth: 最大深度值（mm），默认1200 max_y: 最大Y坐标值（mm），默认150 Returns: tuple: (是否检测到鱼(bool), 满足条件的点数(int)) """ height, width = depth_np.shape fx = calib_params.left_cam.fx fy = calib_params.left_cam.fy cx = calib_params.left_cam.cx cy = calib_params.left_cam.cy # 创建网格坐标 x_grid, y_grid = np.meshgrid(np.arange(width), np.arange(height)) # 计算3D点 Z = depth_np # 深度值（mm） X = (x_grid - cx) * Z / fx Y = (y_grid - cy) * Z / fy # 检测条件： # 1. z+ 方向：深度值在 0-max_depth mm 范围内 # 2. y- 方向：Y坐标小于max_y（相机坐标系中y轴向上为正） # 3. 深度值有效（> 0） valid_depth = (Z > 0) & (Z <= max_depth) valid_y = Y < max_y # 同时满足所有条件的点 fish_mask = valid_depth & valid_y # 统计满足条件的点数 fish_point_count = np.sum(fish_mask) # 如果点数 >= min_points，认为检测到鱼 fish_detected = fish_point_count >= min_points return fish_detected, fish_point_count def downsample_point_cloud(self, points, colors, voxel_size=5.0): """ 使用体素网格对点云进行均匀下采样，去除稀疏点 Args: points: 点云坐标数组 (N, 3) colors: 颜色数组 (N, 3) voxel_size: 体素大小（mm），默认5.0mm Returns: tuple: (下采样后的点云, 下采样后的颜色) """ if len(points) == 0: return points, colors # 计算每个点所属的体素索引 voxel_indices = np.floor(points / voxel_size).astype(np.int32) # 使用字典来存储每个体素中的点（保留第一个点） voxel_dict = {} for i in range(len(points)): voxel_key = tuple(voxel_indices[i]) if voxel_key not in voxel_dict: voxel_dict[voxel_key] = i # 获取保留的点的索引 kept_indices = list(voxel_dict.values()) return points[kept_indices], colors[kept_indices] def filter_point_cloud_kdtree(self, points, colors, radius=5.0, min_neighbors=50): """ 使用KDTree进行半径搜索，过滤掉稀疏点 如果点在指定半径内的邻居数量少于min_neighbors，则移除该点 Args: points: 点云坐标数组 (N, 3) colors: 颜色数组 (N, 3) radius: 搜索半径（mm），默认5.0mm min_neighbors: 最小邻居数量，默认50 Returns: tuple: (过滤后的点云, 过滤后的颜色) """ if len(points) == 0: return points, colors # 构建KDTree tree = cKDTree(points) # 对每个点进行半径搜索（包括点本身） # 使用query_ball_point找到半径内的所有点 neighbor_counts = np.zeros(len(points), dtype=np.int32) # 批量查询以提高效率 for i in range(len(points)): neighbors = tree.query_ball_point(points[i], radius) neighbor_counts[i] = len(neighbors) # 保留邻居数量 >= min_neighbors 的点 valid_mask = neighbor_counts >= min_neighbors filtered_points = points[valid_mask] filtered_colors = colors[valid_mask] return filtered_points, filtered_colors def save_point_cloud(self, left_np, depth_np, calib_params, save_path, base_name): """ 从左图像和深度图生成点云 left_np: 左图像的numpy数组 (H, W, 3) BGR格式 depth_np: 深度图的numpy数组 (H, W) 单位为mm calib_params: 相机标定参数 """ height, width = depth_np.shape fx = calib_params.left_cam.fx fy = calib_params.left_cam.fy cx = calib_params.left_cam.cx cy = calib_params.left_cam.cy # 创建网格坐标 x_grid, y_grid = np.meshgrid(np.arange(width), np.arange(height)) # 计算3D点 Z = depth_np # 深度值（mm） X = (x_grid - cx) * Z / fx Y = (y_grid - cy) * Z / fy # 创建点云 valid_points = Z > 0 # 只使用有效的深度值 points = np.stack([X[valid_points], Y[valid_points], Z[valid_points]], axis=1) # 获取对应的颜色（转换为RGB） colors = left_np[valid_points] colors = colors[:, [2,1,0]] # BGR to RGB # 对点云进行均匀下采样以去除稀疏点 points, colors = self.downsample_point_cloud(points, colors, voxel_size=3.0) #print(f"下采样后点数: {len(points)} (原始点数: {np.sum(valid_points)})") # 使用KDTree过滤稀疏点（半径5mm内少于10个邻居的点将被移除） points, colors = self.filter_point_cloud_kdtree(points, colors, radius=5.0, min_neighbors=10) #print(f"KDTree过滤后点数: {len(points)} (原始点数: {np.sum(valid_points)})") points, colors = self.downsample_point_cloud(points, colors, voxel_size=5.0) print(f"下采样后点数: {len(points)} (原始点数: {np.sum(valid_points)})") # 保存为PLY文件 ply_file = os.path.join(save_path, f"{base_name}_pointcloud.ply") with open(ply_file, 'w') as f: # 写入PLY头部 f.write("ply\n") f.write("format ascii 1.0\n") f.write(f"element vertex {len(points)}\n") f.write("property float x\n") f.write("property float y\n") f.write("property float z\n") f.write("property uchar red\n") f.write("property uchar green\n") f.write("property uchar blue\n") f.write("end_header\n") # 写入点云数据 for i in range(len(points)): f.write(f"{points[i,0]:.3f} {points[i,1]:.3f} {points[i,2]:.3f} " f"{colors[i,0]} {colors[i,1]} {colors[i,2]}\n") print(f"保存点云文件: {ply_file}") def save_images(self, save_path, left_np, right_np, depth_np, frame_count): """保存图像，使用帧计数作为文件名""" timestamp = time.strftime("%Y%m%d_%H%M%S") base_name = f"{timestamp}_{frame_count:06d}" # 保存RGB图像 cv2.imwrite(os.path.join(save_path, f"{base_name}_left.png"), left_np) cv2.imwrite(os.path.join(save_path, f"{base_name}_right.png"), right_np) # 保存原始深度数据（float32格式） depth_file = os.path.join(save_path, f"{base_name}_depth.npy") np.save(depth_file, depth_np) # 获取相机标定参数并生成点云 calibration_params = self.zed.get_camera_information().camera_configuration.calibration_parameters self.save_point_cloud(left_np, depth_np, calibration_params, save_path, base_name) print(f"保存图像、深度数据和点云: {base_name}") def save_fish_point_cloud(self, left_np, depth_np, calib_params, save_path, base_name, max_depth=1000, max_y=100, min_points=500): """ 保存检测到的鱼类点云（只保存满足检测条件的点） left_np: 左图像的numpy数组 (H, W, 3) BGR格式 depth_np: 深度图的numpy数组 (H, W) 单位为mm calib_params: 相机标定参数 save_path: 保存路径 base_name: 文件名基础 max_depth: 最大深度值（mm），默认1000 max_y: 最大Y坐标值（mm），默认100 min_points: 过滤后所需的最少点数，如果少于这个数则不保存，默认500 """ height, width = depth_np.shape fx = calib_params.left_cam.fx fy = calib_params.left_cam.fy cx = calib_params.left_cam.cx cy = calib_params.left_cam.cy # 创建网格坐标 x_grid, y_grid = np.meshgrid(np.arange(width), np.arange(height)) # 计算3D点 Z = depth_np # 深度值（mm） X = (x_grid - cx) * Z / fx Y = (y_grid - cy) * Z / fy # 应用与detect_fish相同的过滤条件 valid_depth = (Z > 0) & (Z <= max_depth) valid_y = Y < max_y fish_mask = valid_depth & valid_y # 只保存满足条件的点 points = np.stack([X[fish_mask], Y[fish_mask], Z[fish_mask]], axis=1) raw_points = points.copy() # 获取对应的颜色（转换为RGB） colors = left_np[fish_mask] colors = colors[:, [2,1,0]] # BGR to RGB raw_colors = colors.copy() # 对点云进行均匀下采样以去除稀疏点 points, colors = self.downsample_point_cloud(points, colors, voxel_size=3.0) # 使用KDTree过滤稀疏点（半径5mm内少于10个邻居的点将被移除） points, colors = self.filter_point_cloud_kdtree(points, colors, radius=5.0, min_neighbors=10) print(f"KDTree过滤后点数: {len(points)} (原始点数: {np.sum(fish_mask)})") points, colors = self.downsample_point_cloud(points, colors, voxel_size=5.0) final_point_count = len(points) print(f"下采样后点数: {final_point_count} (原始点数: {np.sum(fish_mask)})") # 检查过滤后的点数是否满足最小点数要求 if final_point_count < min_points: print(f"跳过保存：过滤后点数 ({final_point_count}) 少于最小要求 ({min_points})") return False # 保存为PLY文件到pointcloud目录 pointcloud_dir = os.path.join(save_path, "pointcloud") self.ensure_save_path(pointcloud_dir) ply_file = os.path.join(pointcloud_dir, f"{base_name}_pointcloud.ply") with open(ply_file, 'w') as f: # 写入PLY头部 f.write("ply\n") f.write("format ascii 1.0\n") f.write(f"element vertex {len(points)}\n") f.write("property float x\n") f.write("property float y\n") f.write("property float z\n") f.write("property uchar red\n") f.write("property uchar green\n") f.write("property uchar blue\n") f.write("end_header\n") # 写入点云数据 for i in range(len(points)): f.write(f"{points[i,0]:.3f} {points[i,1]:.3f} {points[i,2]:.3f} " f"{colors[i,0]} {colors[i,1]} {colors[i,2]}\n") print(f"保存点云文件: {ply_file} (点数: {len(points)})") return True def save_fish_point_cloud_debug(self, left_np, depth_np, calib_params, save_path, base_name, max_depth=1000, max_y=100, min_points=500): """ 在debug模式下保存检测到的鱼类点云（只保存满足检测条件的点） left_np: 左图像的numpy数组 (H, W, 3) BGR格式 depth_np: 深度图的numpy数组 (H, W) 单位为mm calib_params: 相机标定参数 save_path: 保存路径 base_name: 文件名基础 max_depth: 最大深度值（mm），默认1000 max_y: 最大Y坐标值（mm），默认100 min_points: 过滤后所需的最少点数，如果少于这个数则不保存，默认500 """ height, width = depth_np.shape fx = calib_params.left_cam.fx fy = calib_params.left_cam.fy cx = calib_params.left_cam.cx cy = calib_params.left_cam.cy # 创建网格坐标 x_grid, y_grid = np.meshgrid(np.arange(width), np.arange(height)) # 计算3D点 Z = depth_np # 深度值（mm） X = (x_grid - cx) * Z / fx Y = (y_grid - cy) * Z / fy # 应用与detect_fish相同的过滤条件 valid_depth = (Z > 0) & (Z <= max_depth) valid_y = Y < max_y fish_mask = valid_depth & valid_y # 只保存满足条件的点 points = np.stack([X[fish_mask], Y[fish_mask], Z[fish_mask]], axis=1) colors = left_np[fish_mask][:, [2,1,0]] # BGR to RGB if len(points) == 0: print("[DEBUG] 跳过保存：满足过滤条件的原始点数为 0") return False # 对点云进行均匀下采样以去除稀疏点 points, colors = self.downsample_point_cloud(points, colors, voxel_size=3.0) #print(f"[DEBUG] 下采样后点数: {len(points)} (原始点数: {np.sum(fish_mask)})") # 使用KDTree过滤稀疏点（半径5mm内少于50个邻居的点将被移除） points, colors = self.filter_point_cloud_kdtree(points, colors, radius=5.0, min_neighbors=10) print(f"[DEBUG] KDTree过滤后点数: {len(points)} (原始点数: {np.sum(fish_mask)})") points, colors = self.downsample_point_cloud(points, colors, voxel_size=5.0) final_point_count = len(points) print(f"[DEBUG] 下采样后点数: {final_point_count} (原始点数: {np.sum(fish_mask)})") # 保存为PLY文件到pointcloud目录 pointcloud_dir = os.path.join(save_path, "pointcloud") self.ensure_save_path(pointcloud_dir) # 检查过滤后的点数是否满足最小点数要求 if final_point_count < min_points: print(f"[DEBUG] 跳过保存：过滤后点数 ({final_point_count}) 少于最小要求 ({min_points})") return False # 保存为PLY文件 ply_file = os.path.join(pointcloud_dir, f"{base_name}_fish_detected.ply") with open(ply_file, 'w') as f: # 写入PLY头部 f.write("ply\n") f.write("format ascii 1.0\n") f.write(f"element vertex {len(points)}\n") f.write("property float x\n") f.write("property float y\n") f.write("property float z\n") f.write("property uchar red\n") f.write("property uchar green\n") f.write("property uchar blue\n") f.write("end_header\n") # 写入点云数据 for i in range(len(points)): f.write(f"{points[i,0]:.3f} {points[i,1]:.3f} {points[i,2]:.3f} " f"{colors[i,0]} {colors[i,1]} {colors[i,2]}\n") print(f"[DEBUG] 保存检测到的鱼类点云: {ply_file} (点数: {len(points)})") return True def process_frames(self, save_path=None, record=False, save_all=False, debug=False, max_depth=1000, max_y=100, min_points=50, show_yolo_bbox=True): """ 处理帧 save_path: 保存路径 record: 是否录制SVO文件 save_all: 是否保存所有帧 debug: 是否启用debug模式（检测到鱼时保存点云） max_depth: 最大深度值（mm），默认1000 max_y: 最大Y坐标值（mm），默认100 min_points: 检测到鱼所需的最少点数，默认50 show_yolo_bbox: 是否显示YOLO检测的边界框，默认True """ self.save_path = save_path # 保存路径作为类属性 # 确保保存路径存在 if save_path: self.ensure_save_path(save_path) # 如果是录制模式，开始录制 if record and self.camera_mode: if save_path: svo_file = os.path.join(save_path, f"recording_{time.strftime('%Y%m%d_%H%M%S')}.svo") else: svo_file = f"recording_{time.strftime('%Y%m%d_%H%M%S')}.svo" self.start_recording(svo_file) # 获取相机标定参数（用于鱼类检测） calibration_params = self.zed.get_camera_information().camera_configuration.calibration_parameters # 创建运行时参数 runtime_parameters = sl.RuntimeParameters() # 创建图像容器 left_image = sl.Mat() right_image = sl.Mat() depth_map = sl.Mat() # 创建输出目录 if save_path: images_dir = os.path.join(save_path, "images") pointcloud_dir = os.path.join(save_path, "pointcloud") self.ensure_save_path(images_dir) self.ensure_save_path(pointcloud_dir) print(f"图像将保存到: {images_dir}") print(f"点云将保存到: {pointcloud_dir}") frame_count = 0 start_time = time.time() fish_detection_count = 0 frame_skip_interval = 5 # Process every 5th frame # 用于非debug模式的鱼类跟踪 previous_fish_detected = False active_fish_tracking = [] # 列表，每个元素是 (first_detection_frame, fish_number) fish_counter = 0 # 用于生成fish1, fish2等编号 try: while True: # 抓取新的帧 if self.zed.grab(runtime_parameters) == sl.ERROR_CODE.SUCCESS: frame_count += 1 # Skip frames: only process every 5th frame if frame_count % frame_skip_interval != 0: # 跳过帧不处理，直接继续 continue # Skip processing for this frame # Process this frame (every 5th frame) # 获取图像 self.zed.retrieve_image(left_image, sl.VIEW.LEFT) self.zed.retrieve_image(right_image, sl.VIEW.RIGHT) self.zed.retrieve_measure(depth_map, sl.MEASURE.DEPTH) # 转换为numpy数组 left_np = left_image.get_data() right_np = right_image.get_data() depth_np = depth_map.get_data() # YOLO鱼类检测（如果启用） yolo_detected = False yolo_bboxes = [] if self.use_yolo_detector: try: yolo_bboxes = self.detector.detect_fish(left_np) yolo_detected = len(yolo_bboxes) > 0 if yolo_detected: print(f"[帧 {frame_count}] YOLO检测到 {len(yolo_bboxes)} 条鱼") except Exception as e: print(f"[帧 {frame_count}] YOLO检测错误: {e}") # 深度鱼类检测 fish_detected, fish_point_count = self.detect_fish(depth_np, calibration_params, min_points=min_points, max_depth=max_depth, max_y=max_y) # 非debug模式：跟踪鱼类检测并保存点云 if not debug and save_path: # 检测新的鱼类出现（从False到True的转换） # 如果启用了YOLO，需要同时满足YOLO和深度检测；否则只需深度检测 detection_triggered = (self.use_yolo_detector and yolo_detected and fish_detected) or (not self.use_yolo_detector and fish_detected) if detection_triggered and not previous_fish_detected: fish_counter += 1 first_detection_frame = frame_count active_fish_tracking.append((first_detection_frame, fish_counter)) detection_method = "YOLO+深度" if self.use_yolo_detector else "深度" print(f"[帧 {frame_count}] 检测到新鱼类 (fish{fish_counter})！{detection_method}检测，满足条件的点数: {fish_point_count}") # 检查是否有需要保存的点云（first_frame + 12） fish_to_remove = [] for first_frame, fish_num in active_fish_tracking: if frame_count == first_frame + 12: # 如果启用了YOLO，再次检查YOLO是否检测到鱼 should_save = True if self.use_yolo_detector: current_yolo_detected = False try: current_yolo_bboxes = self.detector.detect_fish(left_np) current_yolo_detected = len(current_yolo_bboxes) > 0 except: pass if not current_yolo_detected: should_save = False print(f"[帧 {frame_count}] 跳过保存鱼类点云: fish{fish_num} (YOLO未检测到鱼)") if should_save: # 保存点云 timestamp = time.strftime("%Y%m%d_%H%M%S") base_name = f"{timestamp}_{frame_count:06d}_fish{fish_num}" saved = self.save_fish_point_cloud(left_np, depth_np, calibration_params, save_path, base_name, max_depth=max_depth, max_y=max_y, min_points=min_points) if saved: pointcloud_saved_this_frame = True print(f"[帧 {frame_count}] 保存鱼类点云: fish{fish_num} (首次检测于帧 {first_frame})") else: print(f"[帧 {frame_count}] 跳过保存鱼类点云: fish{fish_num} (过滤后点数不足)") fish_to_remove.append((first_frame, fish_num)) # 移除已保存的鱼类跟踪 for item in fish_to_remove: active_fish_tracking.remove(item) # 如果鱼类消失（从True到False），但保留跟踪直到保存完成 # 这样即使鱼类在保存前消失，我们仍会在first_frame + 12时保存 if not detection_triggered and previous_fish_detected: print(f"[帧 {frame_count}] 鱼类消失，但保留跟踪直到保存完成") # 跟踪当前帧是否成功保存了点云（用于决定是否保存图像） pointcloud_saved_this_frame = False if fish_detected: fish_detection_count += 1 if debug: print(f"[帧 {frame_count}] 检测到鱼类！满足条件的点数: {fish_point_count}") # 如果启用debug模式，保存检测到的鱼类点云 if debug and save_path: # 如果启用了YOLO，只有当YOLO检测到鱼时才保存 should_save_debug = True if self.use_yolo_detector and not yolo_detected: should_save_debug = False print(f"[帧 {frame_count}] [DEBUG] 跳过保存：YOLO未检测到鱼") if should_save_debug: timestamp = time.strftime("%Y%m%d_%H%M%S") base_name = f"{timestamp}_{frame_count:06d}_fish" saved = self.save_fish_point_cloud_debug(left_np, depth_np, calibration_params, save_path, base_name, max_depth=max_depth, max_y=max_y, min_points=min_points) if saved: pointcloud_saved_this_frame = True else: print(f"[帧 {frame_count}] [DEBUG] 跳过保存：过滤后点数不足") # 更新前一次的检测状态 # 如果启用了YOLO，需要同时满足YOLO和深度检测 if self.use_yolo_detector: previous_fish_detected = yolo_detected and fish_detected else: previous_fish_detected = fish_detected # 如果设置保存所有帧 if save_all and save_path: self.save_images(save_path, left_np, right_np, depth_np, frame_count) # 复制左图用于绘制边界框（避免修改原始图像） display_left = left_np.copy() # YOLO检测并绘制边界框（如果启用） if self.use_yolo_detector and show_yolo_bbox and yolo_bboxes: try: display_left = self.detector.draw_bboxes(display_left, yolo_bboxes) except Exception as e: print(f"[帧 {frame_count}] 绘制边界框错误: {e}") # 在图像上显示检测状态和信息 if self.camera_mode: elapsed_time = time.time() - start_time fps = frame_count / elapsed_time if elapsed_time > 0 else 0 status = f"Frame: {frame_count} (processed), FPS: {fps:.1f}" cv2.putText(display_left, status, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2) # 在图像上显示检测状态 if self.use_yolo_detector: detection_status = f"YOLO: {'YES' if yolo_detected else 'NO'} | Depth: {'YES' if fish_detected else 'NO'} ({fish_point_count} pts)" status_color = (0, 255, 0) if (yolo_detected and fish_detected) else (0, 165, 255) if (yolo_detected or fish_detected) else (0, 0, 255) else: detection_status = f"Fish: {'YES' if fish_detected else 'NO'} ({fish_point_count} pts)" status_color = (0, 255, 0) if fish_detected else (0, 0, 255) cv2.putText(display_left, detection_status, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, status_color, 2) # 保存带边界框的图像（只有当点云成功保存时才保存图像） if save_path and pointcloud_saved_this_frame: images_dir = os.path.join(save_path, "images") timestamp = time.strftime("%Y%m%d_%H%M%S") image_filename = f"{timestamp}_{frame_count:06d}_bbox.png" image_path = os.path.join(images_dir, image_filename) cv2.imwrite(image_path, display_left) print(f"[帧 {frame_count}] 保存带边界框的图像: {image_path}") elif not self.camera_mode: # 如果是SVO模式且到达文件末尾 break finally: if record: self.zed.disable_recording() self.close() print(f"总共处理了 {frame_count} 帧") print(f"检测到鱼类的帧数: {fish_detection_count}") if self.camera_mode: elapsed_time = time.time() - start_time print(f"平均帧率: {frame_count/elapsed_time:.1f} FPS") def close(self): # 关闭相机 self.zed.close() def main(): import argparse parser = argparse.ArgumentParser() parser.add_argument('--svo', type=str, help='SVO文件路径', default=None) parser.add_argument('--save_path', type=str, help='保存图像的路径', default='./zed_images') parser.add_argument('--record', action='store_true', help='是否录制新的SVO文件') parser.add_argument('--save_all', action='store_true', help='是否保存所有帧') parser.add_argument('--camera', action='store_true', help='使用实时相机模式') parser.add_argument('--debug', action='store_true', help='启用debug模式：检测到鱼时自动保存点云') parser.add_argument('--max_depth', type=float, help='最大深度值（mm），默认1000', default=1000) parser.add_argument('--max_y', type=float, help='最大Y坐标值（mm），默认100', default=100) parser.add_argument('--min_points', type=int, help='检测到鱼所需的最少点数，默认50', default=500) parser.add_argument('--use_yolo', action='store_true', help='使用YOLO检测器检测鱼（只有YOLO检测到边界框时才保存点云）') parser.add_argument('--yolo_model', type=str, help='YOLO模型路径（.pt文件）', default=None) parser.add_argument('--show_yolo_bbox', action='store_true', default=True, help='显示YOLO检测的边界框（默认启用）') args = parser.parse_args() if args.camera and args.svo: print("错误：不能同时指定相机模式和SVO文件") return # 创建ZED读取器实例 zed_reader = ZEDReader(args.svo, camera_mode=args.camera, use_yolo_detector=args.use_yolo, yolo_model_path=args.yolo_model) # 打开相机或文件 if zed_reader.open(): # 处理帧 zed_reader.process_frames(args.save_path, args.record, args.save_all, args.debug, args.max_depth, args.max_y, args.min_points, args.show_yolo_bbox) if __name__ == "__main__": main()