FishServer/FishMeasure/dataset/zed_reader.py

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()