EG/vr/vr_display.py

"""
VR显示管理模块

此模块负责VR立体渲染和显示，包括：
- 创建左右眼相机
- 设置渲染目标
- 应用畸变校正
- 管理渲染分辨率和质量
"""

from panda3d.core import (
    NodePath, Camera, PerspectiveLens, 
    FrameBufferProperties, WindowProperties, 
    GraphicsOutput, GraphicsEngine,
    Texture, GraphicsBuffer, Vec3, 
    RenderState, ColorWriteAttrib, ShaderAttrib
)
from direct.filter.FilterManager import FilterManager

class VRDisplay:
    """管理VR显示和渲染"""
    
    def __init__(self, vr_system):
        """
        初始化VR显示管理器
        
        Args:
            vr_system: VR系统实例
        """
        self.vr_system = vr_system
        self.world = vr_system.world
        self.device = vr_system.device
        
        # 渲染目标
        self.left_buffer = None
        self.right_buffer = None
        self.left_texture = None
        self.right_texture = None
        
        # 相机
        self.left_camera = None
        self.right_camera = None
        self.left_cam_node = None
        self.right_cam_node = None
        
        # 畸变校正
        self.distortion_manager = None
        self.is_initialized = False
        
        # 渲染设置
        self.render_resolution = vr_system.config.get("render_resolution", [2048, 2048])
        self.ipd = vr_system.config.get("ipd", 0.064)  # 默认瞳距64mm
    
    def initialize(self):
        """初始化VR显示系统"""
        print("正在初始化VR显示系统...")
        
        # 1. 创建渲染目标
        if not self._create_render_targets():
            print("✗ 创建VR渲染目标失败")
            return False
        
        # 2. 创建立体相机
        if not self._setup_stereo_cameras():
            print("✗ 设置VR立体相机失败")
            return False
        
        # 3. 设置畸变校正
        if not self._setup_distortion_correction():
            print("✗ 设置畸变校正失败")
            # 这不是致命错误，可以继续
            print("⚠ 将使用未校正的渲染")
        
        self.is_initialized = True
        print("✓ VR显示系统初始化完成")
        return True
    
    def _create_render_targets(self):
        """创建VR渲染目标"""
        try:
            # 获取图形引擎和窗口
            graphics_engine = self.world.graphicsEngine
            main_window = self.world.win
            
            # 创建帧缓冲属性
            fb_props = FrameBufferProperties()
            fb_props.setRgbColor(True)
            fb_props.setRgbaBits(8, 8, 8, 8)
            fb_props.setDepthBits(24)
            fb_props.setMultisamples(4)  # MSAA 4x
            
            # 创建窗口属性
            win_props = WindowProperties()
            win_props.setSize(self.render_resolution[0], self.render_resolution[1])
            
            # 创建左眼缓冲区
            self.left_buffer = graphics_engine.makeOutput(
                main_window.getPipe(), "left_eye_buffer",
                -1, fb_props, win_props,
                GraphicsOutput.BF_refuse_window,
                main_window.getGsg(), main_window
            )
            
            if not self.left_buffer:
                print("✗ 创建左眼缓冲区失败")
                return False
            
            self.left_buffer.setClearColor((0, 0, 0, 1))
            self.left_texture = Texture()
            self.left_buffer.addRenderTexture(
                self.left_texture, 
                GraphicsOutput.RTM_copy_ram, 
                GraphicsOutput.RTP_color
            )
            
            # 创建右眼缓冲区
            self.right_buffer = graphics_engine.makeOutput(
                main_window.getPipe(), "right_eye_buffer",
                -1, fb_props, win_props,
                GraphicsOutput.BF_refuse_window,
                main_window.getGsg(), main_window
            )
            
            if not self.right_buffer:
                print("✗ 创建右眼缓冲区失败")
                return False
            
            self.right_buffer.setClearColor((0, 0, 0, 1))
            self.right_texture = Texture()
            self.right_buffer.addRenderTexture(
                self.right_texture, 
                GraphicsOutput.RTM_copy_ram, 
                GraphicsOutput.RTP_color
            )
            
            print("✓ VR渲染目标创建成功")
            return True
        except Exception as e:
            print(f"创建VR渲染目标时出错: {e}")
            return False
    
    def _setup_stereo_cameras(self):
        """设置立体相机"""
        try:
            # 获取主相机信息
            main_cam = self.world.camera
            main_cam_pos = main_cam.getPos()
            main_cam_hpr = main_cam.getHpr()
            
            # 创建左眼相机
            self.left_cam_node = Camera("left_eye_camera")
            left_lens = PerspectiveLens()
            left_lens.setFov(90)  # 90度视场角
            left_lens.setNearFar(0.01, 1000)
            self.left_cam_node.setLens(left_lens)
            self.left_camera = NodePath(self.left_cam_node)
            self.left_camera.reparentTo(main_cam)
            
            # 设置左眼相机位置（向左偏移半个瞳距）
            self.left_camera.setPos(-self.ipd/2, 0, 0)
            
            # 创建右眼相机
            self.right_cam_node = Camera("right_eye_camera")
            right_lens = PerspectiveLens()
            right_lens.setFov(90)  # 90度视场角
            right_lens.setNearFar(0.01, 1000)
            self.right_cam_node.setLens(right_lens)
            self.right_camera = NodePath(self.right_cam_node)
            self.right_camera.reparentTo(main_cam)
            
            # 设置右眼相机位置（向右偏移半个瞳距）
            self.right_camera.setPos(self.ipd/2, 0, 0)
            
            # 设置相机显示区域
            self.left_cam_node.setScene(self.world.render)
            self.right_cam_node.setScene(self.world.render)
            
            # 连接相机到缓冲区
            self.left_buffer.setSort(0)
            self.left_buffer.setClearActive(True)
            self.left_buffer.setActive(True)
            self.left_buffer.addDisplayRegion().setCamera(self.left_camera)
            
            self.right_buffer.setSort(0)
            self.right_buffer.setClearActive(True)
            self.right_buffer.setActive(True)
            self.right_buffer.addDisplayRegion().setCamera(self.right_camera)
            
            print("✓ VR立体相机设置成功")
            return True
        except Exception as e:
            print(f"设置VR立体相机时出错: {e}")
            return False
    
    def _setup_distortion_correction(self):
        """设置畸变校正"""
        try:
            # 如果设备提供了畸变校正参数
            if self.device and hasattr(self.device, 'get_distortion_parameters'):
                distortion_params = self.device.get_distortion_parameters()
                
                # 创建畸变校正着色器
                # 这里使用FilterManager来应用后处理效果
                self.distortion_manager = FilterManager(self.world.win, self.world.cam)
                
                # 应用畸变校正着色器到左右眼纹理
                # 具体实现取决于设备提供的畸变参数格式
                
                print("✓ VR畸变校正设置成功")
                return True
            else:
                print("⚠ 设备未提供畸变校正参数")
                return False
        except Exception as e:
            print(f"设置畸变校正时出错: {e}")
            return False
    
    def update(self):
        """更新VR显示（每帧调用）"""
        if not self.is_initialized:
            return
        
        # 1. 获取头部姿态
        if self.device:
            head_pose = self.device.get_head_pose()
            if head_pose:
                self.update_cameras(head_pose)
        
        # 2. 如果使用动态分辨率，根据性能调整
        if self.vr_system.config.get("performance", {}).get("dynamic_resolution", False):
            self._update_dynamic_resolution()
    
    def update_cameras(self, head_pose):
        """
        根据头部姿态更新相机位置
        
        Args:
            head_pose: 包含位置和旋转的头部姿态数据
        """
        if not self.is_initialized:
            return
        
        # 获取主相机
        main_cam = self.world.camera
        
        # 应用头部位置和旋转
        if 'position' in head_pose:
            main_cam.setPos(Vec3(*head_pose['position']))
        
        if 'rotation' in head_pose:
            main_cam.setHpr(Vec3(*head_pose['rotation']))
    
    def render_frame(self):
        """渲染一帧VR画面"""
        if not self.is_initialized:
            return None, None
        
        # 强制渲染左右眼缓冲区
        self.world.graphicsEngine.renderFrame()
        
        # 获取渲染结果
        left_image = self.left_texture.getRamImage()
        right_image = self.right_texture.getRamImage()
        
        return left_image, right_image
    
    def shutdown(self):
        """关闭VR显示系统"""
        if not self.is_initialized:
            return True
        
        print("正在关闭VR显示系统...")
        
        # 移除畸变校正
        if self.distortion_manager:
            self.distortion_manager.cleanup()
            self.distortion_manager = None
        
        # 移除相机
        if self.left_camera:
            self.left_camera.removeNode()
            self.left_camera = None
        
        if self.right_camera:
            self.right_camera.removeNode()
            self.right_camera = None
        
        # 移除缓冲区
        if self.left_buffer:
            self.world.graphicsEngine.removeWindow(self.left_buffer)
            self.left_buffer = None
        
        if self.right_buffer:
            self.world.graphicsEngine.removeWindow(self.right_buffer)
            self.right_buffer = None
        
        self.is_initialized = False
        print("✓ VR显示系统已关闭")
        return True
    
    def set_resolution(self, width, height):
        """
        设置VR渲染分辨率
        
        Args:
            width: 宽度（像素）
            height: 高度（像素）
        """
        if not self.is_initialized:
            print("VR显示系统未初始化，无法设置分辨率")
            return False
        
        try:
            self.render_resolution = [width, height]
            
            # 需要重新创建缓冲区
            self.shutdown()
            self.initialize()
            
            print(f"✓ VR渲染分辨率已设置为 {width}x{height}")
            return True
        except Exception as e:
            print(f"设置VR渲染分辨率时出错: {e}")
            return False
    
    def set_ipd(self, ipd):
        """
        设置瞳距
        
        Args:
            ipd: 瞳距（米）
        """
        if not self.is_initialized:
            print("VR显示系统未初始化，无法设置瞳距")
            return False
        
        try:
            self.ipd = ipd
            
            # 更新相机位置
            if self.left_camera and self.right_camera:
                self.left_camera.setPos(-self.ipd/2, 0, 0)
                self.right_camera.setPos(self.ipd/2, 0, 0)
            
            print(f"✓ VR瞳距已设置为 {ipd}m")
            return True
        except Exception as e:
            print(f"设置VR瞳距时出错: {e}")
            return False
    
    def _update_dynamic_resolution(self):
        """根据性能动态调整渲染分辨率"""
        # 获取当前帧率
        clock = self.world.globalClock
        fps = clock.getAverageFrameRate()
        
        # 目标帧率
        target_fps = self.vr_system.config.get("refresh_rate", 90)
        
        # 如果帧率低于目标的90%，降低分辨率
        if fps < target_fps * 0.9:
            current_width, current_height = self.render_resolution
            new_width = int(current_width * 0.9)
            new_height = int(current_height * 0.9)
            
            # 设置最低分辨率限制
            if new_width >= 1024 and new_height >= 1024:
                self.set_resolution(new_width, new_height)
        
        # 如果帧率高于目标的110%，提高分辨率
        elif fps > target_fps * 1.1:
            current_width, current_height = self.render_resolution
            new_width = int(current_width * 1.1)
            new_height = int(current_height * 1.1)
            
            # 设置最高分辨率限制
            if new_width <= 4096 and new_height <= 4096:
                self.set_resolution(new_width, new_height)