nano_WaveGen/utils/visualize_training.py

#!/usr/bin/env python3
"""
WaveGen 训练结果可视化工具 (独立版本)
自动检索 core_space 目录并可视化训练输出

Usage:
    cd code/WaveGen/nano_WaveGen
    python utils/visualize_training.py
"""

import numpy as np
import viser
import viser.transforms as viser_tf
from typing import Optional, Dict, List, Tuple, Any
import os
from pathlib import Path
import json
import cv2
import time
import webbrowser
from scipy.spatial.transform import Rotation
import threading

# 导入深度转点云模块
try:
    from depth_to_pointcloud import DepthToPointCloud
except ImportError:
    # 尝试从当前目录导入
    import sys
    sys.path.append(str(Path(__file__).parent))
    from depth_to_pointcloud import DepthToPointCloud


class TrainingVisualizer:
    """WaveGen训练结果可视化器"""

    def __init__(self, core_space_dir: str = "core_space", port: int = 8080):
        """
        初始化可视化器

        Args:
            core_space_dir: core_space目录路径（相对于当前工作目录）
            port: 起始端口号（如果占用会自动尝试下一个）
        """
        self.core_space_dir = Path(core_space_dir)
        if not self.core_space_dir.is_absolute():
            self.core_space_dir = Path.cwd() / self.core_space_dir

        # 启动Viser服务器，自动寻找可用端口
        self.server = None
        self.port = port
        max_attempts = 10

        for attempt in range(max_attempts):
            try_port = port + attempt
            try:
                # 不显示默认的config/diagnostics页，直接进入UI
                self.server = viser.ViserServer(port=try_port, show_config=False)
                self.port = try_port
                print(f"🌐 Viser服务器已启动: http://localhost:{try_port}")
                if attempt > 0:
                    print(f"   (端口 {port} 被占用，自动使用端口 {try_port})")
                break
            except OSError as e:
                if "Address already in use" in str(e):
                    if attempt == max_attempts - 1:
                        print(f"❌ 无法找到可用端口 (尝试了 {port}-{try_port})")
                        print(f"   请手动关闭其他实例: pkill -f visualize_training.py")
                        raise
                    continue
                else:
                    raise

        # 可视化句柄
        self.superquadric_handles = []
        self.gt_superquadric_handles = []
        self.camera_handles = []
        self.camera_frustum_handles = []
        self.point_cloud_handle = None
        self.camera_rgb_handle = None
        self.coordinate_frame_handle = None
        self.mesh_handles_pool = {}
        self.object_label_handles = []  # 物体信息标签

        # 当前数据
        self.predictions_npz = None
        self.targets_npz = None
        self.current_sample_path = None
        self.current_frame = 0
        self.original_frame_count = 0
        self.scene_center = np.array([0, 0, 0])
        self.scene_scale = 1.0

        # GUI控件
        self.gui_controls = {}

        # 播放状态
        self.is_playing = False

        # 视频导出状态
        self.is_exporting = False
        self.export_progress = 0
        self.export_camera_pos = None
        self.export_camera_wxyz = None

        # 设置场景
        self.setup_scene()

        # 扫描训练输出
        self.scan_training_outputs()

        # 设置GUI (立即创建，不等待客户端连接)
        self.setup_gui()

        print("✅ 训练可视化器已初始化")
        print(f"📁 监控目录: {self.core_space_dir}")

        if len(self.training_outputs) == 0:
            print("⚠️ 未找到训练输出，请检查 core_space 目录")

    def setup_scene(self):
        """设置场景背景和坐标系"""
        # 设置深蓝色背景（默认）
        self.update_background(wireframe_mode=False)

        # 设置坐标系方向
        self.server.scene.set_up_direction("+y")

    def update_background(self, wireframe_mode: bool):
        """更新场景背景颜色"""
        if wireframe_mode:
            # 线框模式：全黑背景
            bg_color = [0, 0, 0]
        else:
            # 正常模式：深蓝色背景
            bg_color = [13, 13, 38]

        width, height = 1920, 1080
        solid_color_image = np.full((height, width, 3), bg_color, dtype=np.uint8)
        self.server.scene.set_background_image(solid_color_image, format="png")

    def scan_training_outputs(self):
        """扫描core_space目录下的训练输出"""
        self.training_outputs = []

        if not self.core_space_dir.exists():
            print(f"⚠️ core_space目录不存在: {self.core_space_dir}")
            return

        # 查找所有训练输出目录 (格式: YYYYMMDD_HHMMSS_stepN_text2wave)
        # 按时间倒序（最新在前）方便默认选择最新样本
        for output_dir in sorted(self.core_space_dir.glob("*_text2wave"), reverse=True):
            if output_dir.is_dir():
                # 查找样本目录
                sample_dirs = sorted(output_dir.glob("sample_*"))
                if sample_dirs:
                    self.training_outputs.append({
                        'path': output_dir,
                        'name': output_dir.name,
                        'samples': len(sample_dirs)
                    })

        print(f"📦 找到 {len(self.training_outputs)} 个训练输出")
        for output in self.training_outputs:
            print(f"   - {output['name']} ({output['samples']} 样本)")

    def setup_gui(self):
        """设置GUI控件"""
        # 训练输出选择
        with self.server.gui.add_folder("训练输出"):
            if self.training_outputs:
                output_names = [out['name'] for out in self.training_outputs]
                self.gui_controls['output_selector'] = self.server.gui.add_dropdown(
                    "选择训练输出",
                    options=output_names,
                    initial_value=output_names[0]
                )
                self.gui_controls['output_selector'].on_update(self._on_output_change)

                # 样本选择
                self.gui_controls['sample_slider'] = self.server.gui.add_slider(
                    "样本索引",
                    min=0,
                    max=max(0, self.training_outputs[0]['samples'] - 1),
                    step=1,
                    initial_value=0
                )
                self.gui_controls['sample_slider'].on_update(self._on_sample_change)

                self.gui_controls['load_button'] = self.server.gui.add_button("加载样本")
                self.gui_controls['load_button'].on_click(self._on_load_sample)
            else:
                self.server.gui.add_text("状态", initial_value="未找到训练输出")

        # 帧控制
        with self.server.gui.add_folder("帧控制"):
            self.gui_controls['frame_slider'] = self.server.gui.add_slider(
                "当前帧",
                min=0,
                max=23,
                step=1,
                initial_value=0
            )
            self.gui_controls['frame_slider'].on_update(self._on_frame_change)

            self.gui_controls['play_button'] = self.server.gui.add_button("▶ 播放")
            self.gui_controls['play_button'].on_click(self._on_play)

            self.gui_controls['pause_button'] = self.server.gui.add_button("⏸ 暂停")
            self.gui_controls['pause_button'].on_click(self._on_pause)

            self.gui_controls['fps_slider'] = self.server.gui.add_slider(
                "播放FPS",
                min=1,
                max=30,
                step=1,
                initial_value=8
            )

        # 生成结果控制
        with self.server.gui.add_folder("生成结果"):
            self.gui_controls['show_generated'] = self.server.gui.add_checkbox(
                "显示生成的超二次曲面", initial_value=True
            )
            self.gui_controls['show_generated'].on_update(self._on_visibility_change)

            self.gui_controls['generated_opacity'] = self.server.gui.add_slider(
                "生成结果透明度", min=0.1, max=1.0, step=0.05, initial_value=0.7
            )
            self.gui_controls['generated_opacity'].on_update(self._on_opacity_change)

            self.gui_controls['generated_color'] = self.server.gui.add_rgb(
                "生成结果颜色", initial_value=(100, 149, 237)  # 蓝色
            )
            self.gui_controls['generated_color'].on_update(self._on_color_change)

        # Ground Truth控制
        with self.server.gui.add_folder("Ground Truth"):
            self.gui_controls['show_gt'] = self.server.gui.add_checkbox(
                "显示GT超二次曲面", initial_value=True
            )
            self.gui_controls['show_gt'].on_update(self._on_visibility_change)

            self.gui_controls['gt_opacity'] = self.server.gui.add_slider(
                "GT透明度", min=0.1, max=1.0, step=0.05, initial_value=0.5
            )
            self.gui_controls['gt_opacity'].on_update(self._on_opacity_change)

            self.gui_controls['gt_color'] = self.server.gui.add_rgb(
                "GT颜色", initial_value=(255, 99, 71)  # 红色
            )
            self.gui_controls['gt_color'].on_update(self._on_color_change)

            self.gui_controls['show_object_info'] = self.server.gui.add_checkbox(
                "显示物体信息", initial_value=False
            )
            self.gui_controls['show_object_info'].on_update(self._on_visibility_change)

        # 点云控制
        with self.server.gui.add_folder("点云显示"):
            self.gui_controls['show_pointcloud'] = self.server.gui.add_checkbox(
                "显示点云", initial_value=True
            )
            self.gui_controls['show_pointcloud'].on_update(self._on_visibility_change)

            self.gui_controls['pointcloud_size'] = self.server.gui.add_slider(
                "点大小", min=0.001, max=0.02, step=0.001, initial_value=0.008
            )
            self.gui_controls['pointcloud_size'].on_update(self._on_visibility_change)

        # 网格质量
        with self.server.gui.add_folder("渲染设置"):
            self.gui_controls['mesh_resolution'] = self.server.gui.add_slider(
                "网格分辨率", min=10, max=50, step=5, initial_value=25
            )
            self.gui_controls['mesh_resolution'].on_update(self._on_mesh_resolution_change)

            self.gui_controls['show_coordinate'] = self.server.gui.add_checkbox(
                "显示坐标系", initial_value=False
            )
            self.gui_controls['show_coordinate'].on_update(self._on_visibility_change)

            self.gui_controls['wireframe_mode'] = self.server.gui.add_checkbox(
                "线框模式 (黑白边缘)", initial_value=False
            )
            self.gui_controls['wireframe_mode'].on_update(self._on_wireframe_mode_change)

        # 相机控制
        with self.server.gui.add_folder("相机控制"):
            self.gui_controls['reset_view'] = self.server.gui.add_button("重置视角")
            self.gui_controls['reset_view'].on_click(self._on_reset_view)

            self.gui_controls['match_camera'] = self.server.gui.add_button("匹配GT相机")
            self.gui_controls['match_camera'].on_click(self._on_match_camera)

            self.gui_controls['show_target_frustum'] = self.server.gui.add_checkbox(
                "显示GT相机椎体", initial_value=True
            )
            self.gui_controls['show_pred_frustum'] = self.server.gui.add_checkbox(
                "显示预测相机椎体", initial_value=True
            )
            self.gui_controls['show_camera_rgb'] = self.server.gui.add_checkbox(
                "相机视锥显示RGB", initial_value=True
            )
            self.gui_controls['show_target_frustum'].on_update(self._on_visibility_change)
            self.gui_controls['show_pred_frustum'].on_update(self._on_visibility_change)
            self.gui_controls['show_camera_rgb'].on_update(self._on_visibility_change)

        # 视频导出
        with self.server.gui.add_folder("视频导出"):
            self.gui_controls['export_status'] = self.server.gui.add_text(
                "状态", initial_value="就绪"
            )

            self.gui_controls['export_resolution'] = self.server.gui.add_slider(
                "导出分辨率", min=480, max=1080, step=120, initial_value=720
            )

            self.gui_controls['capture_camera_button'] = self.server.gui.add_button(
                "📸 捕获当前视角"
            )
            self.gui_controls['capture_camera_button'].on_click(self._on_capture_camera)

            self.gui_controls['export_viser_button'] = self.server.gui.add_button(
                "💾 导出场景(.viser)"
            )
            self.gui_controls['export_viser_button'].on_click(self._on_export_viser)

            self.gui_controls['export_button'] = self.server.gui.add_button("🎬 导出视频(MP4)")
            self.gui_controls['export_button'].on_click(self._on_export_video)

        print(f"✅ GUI 已设置 - 创建了 {len(self.gui_controls)} 个控件")

    def _on_output_change(self, event):
        """训练输出选择改变"""
        selected_name = event.target.value
        for i, output in enumerate(self.training_outputs):
            if output['name'] == selected_name:
                # 更新样本滑块范围
                max_sample = max(0, output['samples'] - 1)
                self.gui_controls['sample_slider'].max = max_sample
                self.gui_controls['sample_slider'].value = 0
                break

    def _on_sample_change(self, event):
        """样本索引改变"""
        pass  # 用户需要点击"加载样本"按钮

    def _on_load_sample(self, event):
        """加载选中的样本"""
        selected_name = self.gui_controls['output_selector'].value
        sample_idx = int(self.gui_controls['sample_slider'].value)

        # 找到对应的训练输出
        output_path = None
        for output in self.training_outputs:
            if output['name'] == selected_name:
                output_path = output['path']
                break

        if output_path is None:
            print(f"❌ 未找到训练输出: {selected_name}")
            return

        self.load_sample(output_path, sample_idx)

    def load_sample(self, output_path: Path, sample_idx: int):
        """加载样本数据"""
        sample_path = output_path / f"sample_{sample_idx}"

        if not sample_path.exists():
            print(f"❌ 样本目录不存在: {sample_path}")
            return

        print(f"\n{'='*60}")
        print(f"📂 加载样本: {output_path.name}/sample_{sample_idx}")
        print(f"{'='*60}")

        self.current_sample_path = sample_path

        # 加载predictions.npz
        pred_file = sample_path / "predictions.npz"
        if pred_file.exists():
            npz_data = np.load(pred_file, allow_pickle=True)
            self.predictions_npz = {key: npz_data[key] for key in npz_data.files}
            npz_data.close()
            print(f"✅ 加载predictions.npz: {pred_file}")
            if 'frames' in self.predictions_npz:
                print(f"   帧数: {len(self.predictions_npz['frames'])}")
                if 'text' in self.predictions_npz:
                    print(f"   文本: {self.predictions_npz['text']}")
        else:
            self.predictions_npz = None
            print(f"⚠️ 未找到predictions.npz")

        # 加载targets.npz
        target_file = sample_path / "targets.npz"
        if target_file.exists():
            npz_data = np.load(target_file, allow_pickle=True)
            self.targets_npz = {key: npz_data[key] for key in npz_data.files}
            npz_data.close()
            print(f"✅ 加载targets.npz: {target_file}")
            if 'frames' in self.targets_npz:
                print(f"   帧数: {len(self.targets_npz['frames'])}")
                if 'text' in self.targets_npz:
                    print(f"   文本: {self.targets_npz['text']}")
        else:
            self.targets_npz = None
            print(f"⚠️ 未找到targets.npz")

        # 更新帧数
        self.original_frame_count = 0
        if self.predictions_npz and 'frames' in self.predictions_npz:
            self.original_frame_count = len(self.predictions_npz['frames'])
        elif self.targets_npz and 'objects' in self.targets_npz:
            objects = self.targets_npz['objects']
            if hasattr(objects, 'shape') and len(objects.shape) >= 1:
                self.original_frame_count = objects.shape[0]

        if self.original_frame_count > 0:
            self.gui_controls['frame_slider'].max = self.original_frame_count - 1
            self.gui_controls['frame_slider'].value = 0
            self.current_frame = 0
            print(f"📊 总帧数: {self.original_frame_count}")

        # 可视化第一帧
        self.visualize_frame(0)

    def _on_frame_change(self, event):
        """帧滑块改变"""
        frame_idx = int(event.target.value)
        self.visualize_frame(frame_idx)

    def _on_play(self, event):
        """开始播放"""
        self.is_playing = True
        print("▶ 开始播放")

        # 在后台线程播放
        import threading
        threading.Thread(target=self._playback_loop, daemon=True).start()

    def _on_pause(self, event):
        """暂停播放"""
        self.is_playing = False
        print("⏸ 暂停播放")

    def _playback_loop(self):
        """播放循环"""
        while self.is_playing:
            current_frame = int(self.gui_controls['frame_slider'].value)
            next_frame = (current_frame + 1) % self.original_frame_count

            self.gui_controls['frame_slider'].value = next_frame
            self.visualize_frame(next_frame)

            fps = int(self.gui_controls['fps_slider'].value)
            time.sleep(1.0 / fps)

    def _on_visibility_change(self, event):
        """可见性改变"""
        self.visualize_frame(self.current_frame)

    def _on_opacity_change(self, event):
        """透明度改变"""
        self.visualize_frame(self.current_frame)

    def _on_color_change(self, event):
        """颜色改变"""
        self.visualize_frame(self.current_frame)

    def _on_mesh_resolution_change(self, event):
        """网格分辨率改变"""
        # 清空对象池，强制重新生成mesh
        for mesh in self.mesh_handles_pool.values():
            mesh.remove()
        self.mesh_handles_pool.clear()
        self.visualize_frame(self.current_frame)

    def _on_wireframe_mode_change(self, event):
        """线框模式改变"""
        wireframe_mode = event.target.value

        # 更新背景颜色
        self.update_background(wireframe_mode)

        # 清空对象池，强制重新生成mesh（应用线框模式）
        for mesh in self.mesh_handles_pool.values():
            mesh.remove()
        self.mesh_handles_pool.clear()

        # 重新可视化当前帧
        self.visualize_frame(self.current_frame)

    def _on_reset_view(self, event):
        """重置视角"""
        # 设置默认相机位置
        for client in self.server.get_clients().values():
            client.camera.position = (3.0, 2.0, 3.0)
            client.camera.look_at = (0.0, 0.0, 0.0)

    def _on_match_camera(self, event):
        """匹配GT相机视角 (新格式)"""
        if self.targets_npz is None or 'frames' not in self.targets_npz:
            print("⚠️ 没有GT相机数据")
            return

        frame_idx = self.current_frame
        frames = self.targets_npz['frames']

        if frame_idx >= len(frames):
            print("⚠️ 帧索引超出范围")
            return

        frame_data = frames[frame_idx]

        # 转换为字典格式
        if isinstance(frame_data, np.ndarray):
            frame_data = frame_data.item()

        if 'world_info' not in frame_data:
            print("⚠️ 未找到world_info数据")
            return

        world_info = frame_data['world_info']
        camera_position = world_info['camera_position']
        # 存储的是xyzw，需要转成viser的wxyz
        q_xyzw = np.array(world_info['camera_quaternion'], dtype=np.float32)
        wxyz = (float(q_xyzw[3]), float(q_xyzw[0]), float(q_xyzw[1]), float(q_xyzw[2]))

        # 对齐到可视化坐标系：减去scene_center并乘以scene_scale
        cam_pos_vis = (np.array(camera_position) - self.scene_center) * self.scene_scale

        print(f"📷 匹配相机: pos={camera_position}, quat={wxyz}")

        # 设置所有客户端的相机
        for client in self.server.get_clients().values():
            client.camera.position = tuple(cam_pos_vis)
            client.camera.wxyz = wxyz

    def visualize_frame(self, frame_idx: int):
        """可视化指定帧"""
        if self.original_frame_count <= 0:
            return

        frame_idx = int(np.clip(frame_idx, 0, self.original_frame_count - 1))
        self.current_frame = frame_idx

        print(f"\n🎨 可视化帧 {frame_idx}/{self.original_frame_count-1}")

        # 清空旧的可视化
        self.clear_visualization()

        # 获取GUI参数
        show_generated = self.gui_controls['show_generated'].value
        show_gt = self.gui_controls['show_gt'].value
        show_pointcloud = self.gui_controls['show_pointcloud'].value
        show_coordinate = self.gui_controls['show_coordinate'].value

        generated_opacity = self.gui_controls['generated_opacity'].value
        gt_opacity = self.gui_controls['gt_opacity'].value
        generated_color = tuple(self.gui_controls['generated_color'].value)
        gt_color = tuple(self.gui_controls['gt_color'].value)
        mesh_resolution = int(self.gui_controls['mesh_resolution'].value)

        # 提取帧数据
        predictions = self._extract_predictions(frame_idx)
        targets = self._extract_targets(frame_idx)

        # 场景对齐信息：优先使用scene_normalization.json或GT world_info的center/scale；超二次体保持原坐标，点云/相机用该center/scale
        self.scene_center = np.zeros(3, dtype=np.float32)
        self.scene_scale = 1.0

        norm_path = None
        if self.current_sample_path is not None:
            norm_path = self.current_sample_path / "original_data" / "scene_normalization.json"
        loaded_norm = False
        if norm_path is not None and norm_path.exists():
            try:
                with open(norm_path) as f:
                    norm = json.load(f)
                if 'scene_center' in norm:
                    self.scene_center = np.array(norm['scene_center'], dtype=np.float32)
                if 'scene_scale' in norm:
                    self.scene_scale = float(norm['scene_scale'])
                elif 'scene_extent' in norm and norm['scene_extent']:
                    self.scene_scale = 20.0 / float(norm['scene_extent'])
                loaded_norm = True
            except Exception:
                loaded_norm = False

        if not loaded_norm:
            wi = self._get_world_info(frame_idx, source="targets")
            if wi is not None:
                if 'scene_center' in wi:
                    self.scene_center = np.array(wi['scene_center'], dtype=np.float32)
                if 'scene_scale' in wi:
                    try:
                        self.scene_scale = float(wi['scene_scale'])
                    except Exception:
                        pass

        # 线框模式下不显示点云（黑背景下点云不清晰）
        wireframe_mode = self.gui_controls.get('wireframe_mode', None)
        is_wireframe = wireframe_mode.value if wireframe_mode else False

        if show_pointcloud and not is_wireframe:
            # 点云用同一center/scale做归一化
            self._visualize_pointcloud(frame_idx, scene_center=self.scene_center, scene_scale=self.scene_scale)

        # 可视化生成的超二次曲面
        if show_generated and predictions is not None:
            self._visualize_superquadrics(
                predictions,
                color=generated_color,
                opacity=generated_opacity,
                mesh_resolution=mesh_resolution,
                is_gt=False
            )

        # 可视化GT超二次曲面
        if show_gt and targets is not None:
            self._visualize_superquadrics(
                targets,
                color=gt_color,
                opacity=gt_opacity,
                mesh_resolution=mesh_resolution,
                is_gt=True
            )

            # 显示物体信息（如果启用且不在线框模式）
            show_info = self.gui_controls['show_object_info'].value
            if show_info and not is_wireframe:
                self._visualize_object_labels(frame_idx, targets, is_gt=True)

        # 显示坐标系
        if show_coordinate:
            self.coordinate_frame_handle = self.server.scene.add_frame(
                "/coordinate",
                wxyz=(1, 0, 0, 0),
                position=(0, 0, 0),
                axes_length=1.0,
                axes_radius=0.01
            )

        # 可视化相机椎体/RGB（线框模式下不显示）
        if not is_wireframe:
            self._visualize_cameras(frame_idx)

    def _extract_predictions(self, frame_idx: int) -> Optional[np.ndarray]:
        """提取预测数据 (新格式)"""
        if self.predictions_npz is None or 'frames' not in self.predictions_npz:
            return None

        frames = self.predictions_npz['frames']
        if frame_idx >= len(frames):
            return None

        frame_data = frames[frame_idx]

        # 转换为字典格式
        if isinstance(frame_data, np.ndarray):
            frame_data = frame_data.item()

        if 'superquadrics' not in frame_data:
            return None

        superquadrics = frame_data['superquadrics']
        objects_array = []

        for sq in superquadrics:
            # 转换为数组格式 [15 params]
            obj_params = np.zeros(15, dtype=np.float32)
            obj_params[0] = 1.0 if sq['exists'] else 0.0
            obj_params[1:3] = sq['shape']         # epsilon1, epsilon2
            obj_params[3:6] = sq['scale']         # a, b, c
            obj_params[6:9] = sq['translation']   # x, y, z
            obj_params[9:12] = sq['rotation']     # euler angles
            obj_params[12:15] = sq['velocity']    # vx, vy, vz
            objects_array.append(obj_params)

        return np.array(objects_array, dtype=np.float32)

    def _extract_targets(self, frame_idx: int) -> Optional[np.ndarray]:
        """提取GT数据 (新格式)"""
        if self.targets_npz is None or 'frames' not in self.targets_npz:
            return None

        frames = self.targets_npz['frames']
        if frame_idx >= len(frames):
            return None

        frame_data = frames[frame_idx]

        # 转换为字典格式
        if isinstance(frame_data, np.ndarray):
            frame_data = frame_data.item()

        if 'superquadrics' not in frame_data:
            return None

        superquadrics = frame_data['superquadrics']
        objects_array = []

        for sq in superquadrics:
            # 转换为数组格式 [16 params - GT格式包含 inlier_ratio]
            obj_params = np.zeros(16, dtype=np.float32)
            obj_params[0] = 1.0 if sq['exists'] else 0.0
            obj_params[1:3] = sq['shape']              # epsilon1, epsilon2
            obj_params[3:6] = sq['scale']              # a, b, c
            obj_params[6:9] = sq['translation']        # x, y, z
            obj_params[9:12] = sq['rotation']          # euler angles
            obj_params[12] = sq['inlier_ratio']        # inlier ratio (GT specific)
            obj_params[13:16] = sq['velocity']         # vx, vy, vz
            objects_array.append(obj_params)

        return np.array(objects_array, dtype=np.float32)

    def _visualize_superquadrics(self, objects: np.ndarray, color: Tuple,
                                 opacity: float, mesh_resolution: int, is_gt: bool):
        """可视化超二次曲面"""
        prefix = "gt" if is_gt else "gen"
        num_active = 0

        for obj_idx, obj_params in enumerate(objects):
            # 检查存在标志
            if obj_params[0] > 0.5:
                num_active += 1

                try:
                    # 生成mesh
                    vertices, faces = self.generate_superquadric_mesh(
                        obj_params, num_samples=mesh_resolution
                    )

                    # 使用对象池
                    mesh_key = f"{prefix}_{obj_idx}"
                    mesh = self.get_or_create_mesh(
                        mesh_key, vertices, faces, color, opacity
                    )

                    if is_gt:
                        self.gt_superquadric_handles.append(mesh)
                    else:
                        self.superquadric_handles.append(mesh)

                except Exception as e:
                    print(f"❌ 可视化对象{obj_idx}失败: {e}")

        label = "GT" if is_gt else "生成"
        print(f"   {label}对象数: {num_active}")

    def _visualize_object_labels(self, frame_idx: int, objects: np.ndarray, is_gt: bool):
        """在物体上显示信息标签"""
        # 获取原始字典数据以访问inlier_ratio等
        if is_gt and self.targets_npz is not None and 'frames' in self.targets_npz:
            frames = self.targets_npz['frames']
            if frame_idx >= len(frames):
                return

            frame_data = frames[frame_idx]
            if isinstance(frame_data, np.ndarray):
                frame_data = frame_data.item()

            if 'superquadrics' not in frame_data:
                return

            superquadrics = frame_data['superquadrics']

            for obj_idx, sq in enumerate(superquadrics):
                if not sq['exists']:
                    continue

                # 获取物体位置（用于放置标签）
                translation = sq['translation']
                scale = sq['scale']

                # 标签位置：物体中心上方
                label_position = (
                    float(translation[0]),
                    float(translation[1]) + float(scale[1]) * 1.5,  # 在物体上方
                    float(translation[2])
                )

                # 构建信息文本
                inlier_ratio = sq.get('inlier_ratio', 0.0)
                shape = sq.get('shape', [0, 0])

                info_text = (
                    f"ID: {obj_idx}\n"
                    f"Density: {inlier_ratio:.3f}\n"
                    f"Shape: ε1={shape[0]:.2f}, ε2={shape[1]:.2f}\n"
                    f"Size: {scale[0]:.2f}×{scale[1]:.2f}×{scale[2]:.2f}"
                )

                # 添加文本标签
                # 使用时间戳确保名称唯一，避免冲突
                label_name = f"/object_label_f{frame_idx}_o{obj_idx}"
                try:
                    label_handle = self.server.scene.add_label(
                        label_name,
                        text=info_text,
                        position=label_position
                    )
                    self.object_label_handles.append(label_handle)
                except Exception as e:
                    print(f"⚠️ 创建标签失败: {e}")

    def _visualize_pointcloud(self, frame_idx: int, scene_center: Optional[np.ndarray] = None, scene_scale: Optional[float] = None):
        """可视化点云"""
        if self.current_sample_path is None:
            return

        # 查找原始数据
        original_data_dir = self.current_sample_path / "original_data"
        if not original_data_dir.exists():
            print("⚠️ 未找到original_data目录")
            return

        # 加载深度图和RGB
        depth_file = self._find_depth_file(original_data_dir, frame_idx)
        rgb_file = original_data_dir / "rgb" / f"frame_{frame_idx:03d}.png"

        if depth_file is None or not rgb_file.exists():
            print(f"⚠️ 未找到帧{frame_idx}的深度图或RGB")
            return

        try:
            # 加载数据
            depth = self._load_depth(depth_file, frame_idx)
            if depth.ndim == 2:
                depth = depth[:, :, None]
            rgb = self._load_rgb(rgb_file)

            # 加载相机内参
            camera_K = None
            metadata_file = original_data_dir / "metadata.json"
            if metadata_file.exists():
                with open(metadata_file) as f:
                    metadata = json.load(f)
                if 'camera' in metadata and 'K' in metadata['camera']:
                    camera_K = np.array(metadata['camera']['K'], dtype=np.float32)
            if camera_K is None:
                h, w = depth.shape[:2]
                camera_K = np.array([[w, 0, w/2], [0, h, h/2], [0, 0, 1]], dtype=np.float32)

            # 使用GT world_info作为相机位姿
            world_info = self._get_world_info(frame_idx, source="targets")
            camera_position = np.zeros(3, dtype=np.float32)
            camera_quat_xyzw = np.array([0, 0, 0, 1], dtype=np.float32)
            if world_info is not None and 'camera_position' in world_info:
                camera_position = np.array(world_info['camera_position'], dtype=np.float32)
                if 'camera_quaternion' in world_info:
                    # 存储为xyzw，直接使用
                    camera_quat_xyzw = np.array(world_info['camera_quaternion'], dtype=np.float32)

            # 转换为点云（标准化到[-10,10]），使用提供的center/scale以与训练对齐
            converter = DepthToPointCloud()
            _, points_norm, _, depth_center, depth_extent = converter.depth_to_normalized_pointcloud_movi(
                depth=depth,
                segmentation=None,
                camera_K=camera_K,
                camera_position=camera_position,
                camera_quaternion=camera_quat_xyzw,
                resolution=depth.shape[0],
                convert_to_zdepth=True,
                scene_center_override=scene_center,
                scene_scale_override=scene_scale
            )

            valid_mask = depth[:, :, 0] > 0
            points = points_norm[valid_mask]
            colors = rgb.reshape(-1, 3)[valid_mask.reshape(-1)]

            # 更新场景中心和尺度（供相机椎体使用）
            if scene_center is not None and scene_scale is not None:
                self.scene_center = np.array(scene_center, dtype=np.float32)
                self.scene_scale = float(scene_scale)
            else:
                self.scene_center = depth_center
                self.scene_scale = 20.0 / max(depth_extent, 1e-6)

            # 显示点云
            point_size = self.gui_controls['pointcloud_size'].value
            self.point_cloud_handle = self.server.scene.add_point_cloud(
                "/pointcloud",
                points=points,
                colors=colors,
                point_size=point_size
            )

            print(f"   点云: {len(points)} 个点")

        except Exception as e:
            print(f"❌ 加载点云失败: {e}")

    def _find_depth_file(self, original_data_dir: Path, frame_idx: int) -> Optional[Path]:
        """查找深度文件（支持合并的npz和单独的npy）"""
        depth_dir = original_data_dir / "depth"
        if not depth_dir.exists():
            return None

        # 检查合并的npz
        merged_npz = depth_dir / "depth_merge.npz"
        if merged_npz.exists():
            return merged_npz

        # 检查单独的npy
        npy_file = depth_dir / f"frame_{frame_idx:03d}.npy"
        if npy_file.exists():
            return npy_file

        return None

    def _load_depth(self, depth_file: Path, frame_idx: int) -> np.ndarray:
        """加载深度数据"""
        if depth_file.suffix == '.npz':
            # 从合并的npz加载
            data = np.load(depth_file)
            frame_key = f"frame_{frame_idx:03d}"
            return data[frame_key]
        else:
            # 从单独的npy加载
            return np.load(depth_file)

    def _load_rgb(self, rgb_path: Path) -> np.ndarray:
        """加载RGB图像"""
        img = cv2.imread(str(rgb_path))
        if img is None:
            raise FileNotFoundError(f"Failed to load RGB image: {rgb_path}")
        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    def _get_world_info(self, frame_idx: int, source: str = "targets") -> Optional[Dict[str, np.ndarray]]:
        """从pred/target获取世界/相机信息"""
        data = self.targets_npz if source == "targets" else self.predictions_npz
        if data is None:
            return None

        if 'frames' in data:
            frames = data['frames']
            if frame_idx < len(frames):
                entry = frames[frame_idx]
                if hasattr(entry, 'item'):
                    try:
                        entry = entry.item()
                    except Exception:
                        pass
                if isinstance(entry, dict) and 'world_info' in entry:
                    return entry['world_info']

        # 兼容旧格式 world 张量
        if 'world' in data:
            world = data['world']
            if hasattr(world, 'shape') and world.shape[0] > frame_idx and world.shape[-1] >= 7:
                wp = world[frame_idx]
                scene_center = world[frame_idx, 8:11] if world.shape[-1] >= 11 else np.zeros(3, dtype=np.float32)
                return {
                    'camera_position': wp[:3],
                    'camera_quaternion': wp[3:7],
                    'scene_scale': float(wp[7]) if len(wp) > 7 else 1.0,
                    'scene_center': scene_center,
                }
        return None

    def _visualize_cameras(self, frame_idx: int):
        """可视化相机椎体与RGB"""
        for h in self.camera_frustum_handles:
            h.remove()
        self.camera_frustum_handles = []
        if self.camera_rgb_handle is not None:
            self.camera_rgb_handle.remove()
            self.camera_rgb_handle = None

        show_target = self.gui_controls.get('show_target_frustum', None)
        show_pred = self.gui_controls.get('show_pred_frustum', None)
        show_rgb = self.gui_controls.get('show_camera_rgb', None)
        if show_target is None or show_pred is None or show_rgb is None:
            return
        if not (show_target.value or show_pred.value):
            return

        original_data_dir = None
        rgb_image = None
        if show_rgb.value and self.current_sample_path is not None:
            original_data_dir = self.current_sample_path / "original_data"
            if original_data_dir.exists():
                rgb_path = original_data_dir / "rgb" / f"frame_{frame_idx:03d}.png"
                if rgb_path.exists():
                    try:
                        rgb_image = self._load_rgb(rgb_path)
                    except Exception:
                        rgb_image = None

        # FOV估计
        fov = np.deg2rad(60.0)
        aspect = 1.0
        if rgb_image is not None:
            h, w = rgb_image.shape[:2]
            aspect = w / max(h, 1)
            metadata_file = (self.current_sample_path / "original_data" / "metadata.json") if self.current_sample_path else None
            fx = None
            if metadata_file and metadata_file.exists():
                try:
                    with open(metadata_file) as f:
                        metadata = json.load(f)
                    if 'camera' in metadata and 'K' in metadata['camera']:
                        K = np.array(metadata['camera']['K'], dtype=np.float32)
                        fx = K[0, 0]
                except Exception:
                    fx = None
            if fx is not None and w > 0:
                fov = 2 * np.arctan(w / (2 * fx))

        def add_frustum(world_info: Dict, name: str, color: Tuple[int, int, int]):
            if world_info is None:
                return
            cam_pos = np.array(world_info.get('camera_position', np.zeros(3)), dtype=np.float32)
            cam_quat = np.array(world_info.get('camera_quaternion', [0, 0, 0, 1]), dtype=np.float32)  # xyzw
            if cam_quat.shape[0] == 4:
                wxyz = (float(cam_quat[3]), float(cam_quat[0]), float(cam_quat[1]), float(cam_quat[2]))
            else:
                wxyz = (1.0, 0.0, 0.0, 0.0)

            # 将位置移到可视化坐标系（减去中心再按场景尺度缩放）
            pos = (cam_pos - self.scene_center) * getattr(self, "scene_scale", 1.0)

            frustum = self.server.scene.add_camera_frustum(
                f"/{name}",
                fov=fov,
                aspect=aspect,
                scale=2.0,
                wxyz=wxyz,
                position=pos,
                image=rgb_image if show_rgb.value else None,
                color=tuple(int(c) for c in color)
            )
            self.camera_frustum_handles.append(frustum)

        if show_pred.value:
            add_frustum(self._get_world_info(frame_idx, source="predictions"), "pred_camera_frustum", (100, 149, 237))
        if show_target.value:
            add_frustum(self._get_world_info(frame_idx, source="targets"), "gt_camera_frustum", (255, 99, 71))

    def generate_superquadric_mesh(self, params, num_samples=25):
        """生成超二次曲面mesh"""
        # 解析参数
        epsilon = [params[1], params[2]]
        scale = [params[3], params[4], params[5]]
        translation = [params[6], params[7], params[8]]
        rotation = [params[9], params[10], params[11]] if len(params) >= 12 else [0, 0, 0]

        # 生成参数网格
        eta = np.linspace(-np.pi/2, np.pi/2, num_samples)
        omega = np.linspace(-np.pi, np.pi, num_samples)

        vertices = []
        faces = []

        # 生成旋转矩阵
        rot = Rotation.from_euler('ZYX', rotation)
        rot_matrix = rot.as_matrix()

        # 生成顶点
        for i, e in enumerate(eta):
            for j, w in enumerate(omega):
                # 超二次曲面参数方程
                cos_eta = np.sign(np.cos(e)) * np.abs(np.cos(e))**epsilon[0]
                sin_eta = np.sign(np.sin(e)) * np.abs(np.sin(e))**epsilon[0]
                cos_omega = np.sign(np.cos(w)) * np.abs(np.cos(w))**epsilon[1]
                sin_omega = np.sign(np.sin(w)) * np.abs(np.sin(w))**epsilon[1]

                # 局部坐标
                x_local = scale[0] * cos_eta * cos_omega
                y_local = scale[1] * cos_eta * sin_omega
                z_local = scale[2] * sin_eta

                # 应用旋转和平移
                point_local = np.array([x_local, y_local, z_local])
                point_global = rot_matrix @ point_local + np.array(translation)

                vertices.append(point_global)

        vertices = np.array(vertices)

        # 生成面片
        for i in range(num_samples - 1):
            for j in range(num_samples - 1):
                idx1 = i * num_samples + j
                idx2 = i * num_samples + (j + 1) % num_samples
                idx3 = (i + 1) * num_samples + j
                idx4 = (i + 1) * num_samples + (j + 1) % num_samples

                faces.append([idx1, idx2, idx3])
                faces.append([idx2, idx4, idx3])

        return vertices, np.array(faces)

    def get_or_create_mesh(self, key: str, vertices, faces, color, opacity):
        """获取或创建mesh（对象池）"""
        # 检查是否启用线框模式
        wireframe_mode = self.gui_controls.get('wireframe_mode', None)
        is_wireframe = wireframe_mode.value if wireframe_mode else False

        # 线框模式：强制白色，完全不透明
        if is_wireframe:
            display_color = (255, 255, 255)
            display_opacity = 1.0
        else:
            display_color = color
            display_opacity = opacity

        if key in self.mesh_handles_pool:
            mesh = self.mesh_handles_pool[key]
            mesh.vertices = vertices
            mesh.vertex_colors = None
            mesh.wireframe = is_wireframe
            mesh.opacity = display_opacity
            mesh.visible = True

            # 更新颜色
            color_array = np.array(display_color, dtype=np.uint8)
            if color_array.max() <= 1.0:
                color_array = (color_array * 255).astype(np.uint8)
            mesh.color = tuple(color_array)
        else:
            # 创建新mesh
            color_array = np.array(display_color, dtype=np.uint8)
            if color_array.max() <= 1.0:
                color_array = (color_array * 255).astype(np.uint8)

            mesh = self.server.scene.add_mesh_simple(
                name=f"/mesh_{key}",
                vertices=vertices,
                faces=faces,
                color=tuple(color_array),
                opacity=display_opacity,
                wireframe=is_wireframe,
                flat_shading=False
            )
            self.mesh_handles_pool[key] = mesh

        return mesh

    def clear_visualization(self):
        """清空可视化"""
        # 隐藏所有mesh
        for mesh in self.mesh_handles_pool.values():
            mesh.visible = False

        # 清空句柄列表
        self.superquadric_handles = []
        self.gt_superquadric_handles = []

        # 删除点云
        if self.point_cloud_handle is not None:
            self.point_cloud_handle.remove()
            self.point_cloud_handle = None

        # 删除相机椎体/RGB
        for handle in self.camera_frustum_handles:
            handle.remove()
        self.camera_frustum_handles = []
        if self.camera_rgb_handle is not None:
            self.camera_rgb_handle.remove()
            self.camera_rgb_handle = None

        # 删除坐标系
        if self.coordinate_frame_handle is not None:
            self.coordinate_frame_handle.remove()
            self.coordinate_frame_handle = None

        # 删除物体信息标签
        for handle in self.object_label_handles:
            try:
                handle.remove()
            except (KeyError, AttributeError):
                # 标签可能已经被删除，忽略错误
                pass
        self.object_label_handles = []

    def _on_capture_camera(self, event):
        """捕获当前相机视角"""
        clients = list(self.server.get_clients().values())
        if not clients:
            print("⚠️ 没有连接的客户端")
            self.gui_controls['export_status'].value = "错误: 没有连接的客户端"
            return

        # 获取第一个客户端的相机参数
        client = clients[0]
        self.export_camera_pos = np.array(client.camera.position)
        self.export_camera_wxyz = np.array(client.camera.wxyz)

        print(f"📸 已捕获相机视角: pos={self.export_camera_pos}, wxyz={self.export_camera_wxyz}")
        self.gui_controls['export_status'].value = f"已捕获视角: {self.export_camera_pos}"

    def _on_export_viser(self, event):
        """导出为viser场景文件（可交互）"""
        if self.current_sample_path is None:
            print("⚠️ 请先加载样本")
            self.gui_controls['export_status'].value = "错误: 请先加载样本"
            return

        if self.original_frame_count <= 0:
            print("⚠️ 没有帧可以导出")
            self.gui_controls['export_status'].value = "错误: 没有帧可以导出"
            return

        # 在后台线程导出
        threading.Thread(target=self._export_viser_thread, daemon=True).start()

    def _export_viser_thread(self):
        """导出viser场景文件（带动画）"""
        try:
            print(f"\n{'='*60}")
            print(f"💾 开始导出Viser场景")
            print(f"{'='*60}")

            # 获取当前客户端的相机参数
            clients = list(self.server.get_clients().values())
            camera_params = None
            if clients:
                client = clients[0]
                cam_pos = client.camera.position
                cam_lookat = client.camera.look_at
                cam_up = client.camera.up_direction

                # 生成viser URL参数格式
                camera_params = (
                    f"&initialCameraPosition={cam_pos[0]:.3f},{cam_pos[1]:.3f},{cam_pos[2]:.3f}"
                    f"&initialCameraLookAt={cam_lookat[0]:.3f},{cam_lookat[1]:.3f},{cam_lookat[2]:.3f}"
                    f"&initialCameraUp={cam_up[0]:.3f},{cam_up[1]:.3f},{cam_up[2]:.3f}"
                )
                print(f"   📸 记录相机视角:")
                print(f"      位置: {cam_pos}")
                print(f"      朝向: {cam_lookat}")
                print(f"      向上: {cam_up}")

            # 获取FPS
            fps = int(self.gui_controls['fps_slider'].value)

            # 创建输出目录
            output_dir = self.core_space_dir / "exports"
            output_dir.mkdir(exist_ok=True)

            # 生成文件名
            selected_output = self.gui_controls['output_selector'].value
            sample_idx = int(self.gui_controls['sample_slider'].value)
            step_info = "unknown"
            if "step" in selected_output:
                try:
                    step_part = selected_output.split("_step")[1].split("_")[0]
                    step_info = f"step{step_part}"
                except:
                    pass

            timestamp = time.strftime("%Y%m%d_%H%M%S")
            experiment_name = selected_output.split("_")[0]
            output_file = output_dir / f"{experiment_name}_{step_info}_sample{sample_idx}_{timestamp}.viser"

            print(f"   输出文件: {output_file}")
            print(f"   帧数: {self.original_frame_count}")
            print(f"   FPS: {fps}")

            # 获取场景序列化器
            serializer = self.server.get_scene_serializer()

            # 记录初始状态（第一帧）
            self.visualize_frame(0)
            serializer.insert_sleep(1.0 / fps)

            # 逐帧更新并记录
            for frame_idx in range(1, self.original_frame_count):
                self.export_progress = int((frame_idx + 1) / self.original_frame_count * 100)
                self.gui_controls['export_status'].value = f"导出中... {self.export_progress}%"

                # 更新场景（会自动更新viser场景）
                self.visualize_frame(frame_idx)

                # 添加帧延迟
                serializer.insert_sleep(1.0 / fps)

                print(f"   记录帧 {frame_idx+1}/{self.original_frame_count}")

            # 序列化并保存
            data = serializer.serialize()
            output_file.write_bytes(data)

            print(f"✅ 场景导出完成: {output_file}")
            print(f"   文件大小: {len(data) / 1024 / 1024:.2f} MB")
            print(f"\n📖 查看方式:")
            print(f"   1. 安装viser客户端: viser-build-client --output-dir viser-client/")
            print(f"   2. 启动HTTP服务器: python -m http.server 8000")

            # 生成完整URL（带相机参数）
            base_url = f"http://localhost:8000/viser-client/?playbackPath=http://localhost:8000/exports/{output_file.name}"
            if camera_params:
                full_url = base_url + camera_params
                print(f"   3. 打开浏览器（带相机视角）:")
                print(f"      {full_url}")
            else:
                print(f"   3. 打开浏览器:")
                print(f"      {base_url}")

            relative_path = output_file.relative_to(self.core_space_dir)
            self.gui_controls['export_status'].value = f"完成! {relative_path}"

            # 提供下载
            clients = list(self.server.get_clients().values())
            if clients:
                clients[0].send_file_download(output_file.name, data)
                print(f"   💾 已发送下载到浏览器")

        except Exception as e:
            print(f"❌ 导出失败: {e}")
            import traceback
            traceback.print_exc()
            self.gui_controls['export_status'].value = f"错误: {str(e)}"

    def _on_export_video(self, event):
        """导出视频"""
        if self.is_exporting:
            print("⚠️ 正在导出中，请等待...")
            return

        if self.current_sample_path is None:
            print("⚠️ 请先加载样本")
            self.gui_controls['export_status'].value = "错误: 请先加载样本"
            return

        if self.original_frame_count <= 0:
            print("⚠️ 没有帧可以导出")
            self.gui_controls['export_status'].value = "错误: 没有帧可以导出"
            return

        # 检查是否有连接的客户端
        clients = list(self.server.get_clients().values())
        if not clients:
            print("⚠️ 没有连接的客户端")
            self.gui_controls['export_status'].value = "错误: 请先在浏览器中打开viser界面"
            return

        # 每次导出都获取最新的相机视角（重要！）
        # 无论之前是否捕获过，都使用当前最新的视角
        client = clients[0]
        self.export_camera_pos = np.array(client.camera.position)
        self.export_camera_wxyz = np.array(client.camera.wxyz)
        print(f"📸 使用当前视角: pos={self.export_camera_pos}, wxyz={self.export_camera_wxyz}")

        # 在后台线程导出视频
        threading.Thread(target=self._export_video_thread_screenshot, daemon=True).start()

    def _export_video_thread(self):
        """视频导出线程"""
        try:
            self.is_exporting = True
            self.gui_controls['export_status'].value = "正在导出..."

            # 确保场景归一化参数已设置（通过可视化当前帧来初始化）
            if not hasattr(self, 'scene_center') or self.scene_center is None:
                print("   初始化场景参数...")
                self.visualize_frame(self.current_frame)

            # 获取参数
            fps = int(self.gui_controls['fps_slider'].value)
            resolution = int(self.gui_controls['export_resolution'].value)

            # 创建输出目录 - 放在core_space根目录下
            output_dir = self.core_space_dir / "exports"
            output_dir.mkdir(exist_ok=True)

            # 提取实验信息
            selected_output = self.gui_controls['output_selector'].value
            sample_idx = int(self.gui_controls['sample_slider'].value)

            # 从输出名称提取步数 (例如: 20251205_184253_step5_text2wave -> step5)
            step_info = "unknown"
            if "step" in selected_output:
                try:
                    step_part = selected_output.split("_step")[1].split("_")[0]
                    step_info = f"step{step_part}"
                except:
                    pass

            # 生成输出文件名: {实验名}_{step}_sample{idx}_{timestamp}.mp4
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            experiment_name = selected_output.split("_")[0]  # 取日期部分作为实验名
            output_file = output_dir / f"{experiment_name}_{step_info}_sample{sample_idx}_{timestamp}.mp4"

            print(f"\n{'='*60}")
            print(f"🎬 开始导出视频")
            print(f"{'='*60}")
            print(f"   实验: {selected_output}")
            print(f"   样本: {sample_idx}")
            print(f"   输出文件: {output_file}")
            print(f"   帧数: {self.original_frame_count}")
            print(f"   FPS: {fps}")
            print(f"   分辨率: {resolution}x{resolution}")
            print(f"   相机位置: {self.export_camera_pos}")
            print(f"   相机旋转: {self.export_camera_wxyz}")

            # 尝试使用imageio（更好的兼容性），如果不可用则使用OpenCV
            try:
                import imageio
                use_imageio = True
                print("   使用 imageio 进行视频编码（H.264）")
            except ImportError:
                use_imageio = False
                print("   使用 OpenCV 进行视频编码")

            if use_imageio:
                # 使用imageio-ffmpeg，生成高兼容性的H.264视频
                # 注意：必须指定format='FFMPEG'来确保使用FFmpeg插件
                writer = imageio.get_writer(
                    str(output_file),
                    format='FFMPEG',
                    mode='I',
                    fps=fps,
                    codec='libx264',
                    pixelformat='yuv420p',  # 确保兼容性
                    output_params=['-crf', '18']  # H.264质量参数，18是高质量
                )

                # 渲染每一帧
                for frame_idx in range(self.original_frame_count):
                    self.export_progress = int((frame_idx + 1) / self.original_frame_count * 100)
                    self.gui_controls['export_status'].value = f"导出中... {self.export_progress}%"

                    # 渲染帧
                    frame_image = self._render_frame_offline(
                        frame_idx,
                        resolution=resolution,
                        camera_pos=self.export_camera_pos,
                        camera_wxyz=self.export_camera_wxyz
                    )

                    # 写入视频（imageio需要RGB格式）
                    if frame_image is not None:
                        writer.append_data(frame_image)

                    print(f"   渲染帧 {frame_idx+1}/{self.original_frame_count}")

                writer.close()

            else:
                # 使用OpenCV，尝试更兼容的编码器
                # 尝试顺序: H264 -> avc1 -> X264 -> mp4v
                codecs_to_try = [
                    ('H264', 'H.264'),
                    ('avc1', 'H.264 (AVC1)'),
                    ('X264', 'X264'),
                    ('mp4v', 'MPEG-4')
                ]

                writer = None
                used_codec = None

                for codec_fourcc, codec_name in codecs_to_try:
                    try:
                        fourcc = cv2.VideoWriter_fourcc(*codec_fourcc)
                        test_writer = cv2.VideoWriter(
                            str(output_file),
                            fourcc,
                            fps,
                            (resolution, resolution)
                        )
                        if test_writer.isOpened():
                            writer = test_writer
                            used_codec = codec_name
                            print(f"   使用编码器: {codec_name}")
                            break
                        else:
                            test_writer.release()
                    except:
                        continue

                if writer is None:
                    raise RuntimeError("无法初始化视频编码器")

                # 渲染每一帧
                for frame_idx in range(self.original_frame_count):
                    self.export_progress = int((frame_idx + 1) / self.original_frame_count * 100)
                    self.gui_controls['export_status'].value = f"导出中... {self.export_progress}%"

                    # 渲染帧
                    frame_image = self._render_frame_offline(
                        frame_idx,
                        resolution=resolution,
                        camera_pos=self.export_camera_pos,
                        camera_wxyz=self.export_camera_wxyz
                    )

                    # 写入视频（OpenCV需要BGR格式）
                    if frame_image is not None:
                        writer.write(cv2.cvtColor(frame_image, cv2.COLOR_RGB2BGR))

                    print(f"   渲染帧 {frame_idx+1}/{self.original_frame_count}")

                writer.release()

            print(f"✅ 视频导出完成: {output_file}")
            relative_path = output_file.relative_to(self.core_space_dir)
            self.gui_controls['export_status'].value = f"完成! {relative_path}"

        except Exception as e:
            print(f"❌ 导出视频失败: {e}")
            import traceback
            traceback.print_exc()
            self.gui_controls['export_status'].value = f"错误: {str(e)}"

        finally:
            self.is_exporting = False

    def _export_video_thread_screenshot(self):
        """视频导出线程（基于截图viser界面）"""
        try:
            self.is_exporting = True
            self.gui_controls['export_status'].value = "正在导出..."

            # 获取参数
            fps = int(self.gui_controls['fps_slider'].value)

            # 创建输出目录
            output_dir = self.core_space_dir / "exports"
            output_dir.mkdir(exist_ok=True)

            # 提取实验信息并生成文件名
            selected_output = self.gui_controls['output_selector'].value
            sample_idx = int(self.gui_controls['sample_slider'].value)
            step_info = "unknown"
            if "step" in selected_output:
                try:
                    step_part = selected_output.split("_step")[1].split("_")[0]
                    step_info = f"step{step_part}"
                except:
                    pass

            timestamp = time.strftime("%Y%m%d_%H%M%S")
            experiment_name = selected_output.split("_")[0]
            output_file = output_dir / f"{experiment_name}_{step_info}_sample{sample_idx}_{timestamp}.mp4"

            print(f"\n{'='*60}")
            print(f"🎬 开始导出视频（截图模式）")
            print(f"{'='*60}")
            print(f"   实验: {selected_output}")
            print(f"   样本: {sample_idx}")
            print(f"   输出文件: {output_file}")
            print(f"   帧数: {self.original_frame_count}")
            print(f"   FPS: {fps}")
            print(f"   方法: 直接截取Viser显示画面")

            # 检查selenium
            try:
                from selenium import webdriver
                from selenium.webdriver.chrome.options import Options
                from selenium.webdriver.common.by import By
                import time as time_module
                use_selenium = True
                print("   ✅ 使用 Selenium 截图")
            except ImportError:
                print("   ⚠️  Selenium未安装，使用逐帧渲染方法")
                print("   提示: pip install selenium")
                use_selenium = False

            if use_selenium:
                # 使用Selenium截图方法
                frames = []

                # 配置Chrome
                chrome_options = Options()
                chrome_options.add_argument('--headless')  # 无头模式
                chrome_options.add_argument('--no-sandbox')
                chrome_options.add_argument('--disable-dev-shm-usage')
                chrome_options.add_argument('--window-size=1920,1080')

                try:
                    driver = webdriver.Chrome(options=chrome_options)
                    url = f"http://localhost:{self.port}"
                    driver.get(url)
                    print(f"   📱 打开浏览器: {url}")

                    # 等待页面加载
                    time_module.sleep(3)

                    # 逐帧截图
                    for frame_idx in range(self.original_frame_count):
                        self.export_progress = int((frame_idx + 1) / self.original_frame_count * 100)
                        self.gui_controls['export_status'].value = f"截图中... {self.export_progress}%"

                        # 通过GUI更新帧
                        self.gui_controls['frame_slider'].value = frame_idx
                        time_module.sleep(0.3)  # 等待渲染

                        # 截图
                        screenshot = driver.get_screenshot_as_png()
                        img = cv2.imdecode(np.frombuffer(screenshot, np.uint8), cv2.IMREAD_COLOR)
                        frames.append(img)

                        print(f"   截图帧 {frame_idx+1}/{self.original_frame_count}")

                    driver.quit()

                    # 使用imageio写入视频
                    try:
                        import imageio
                        writer = imageio.get_writer(
                            str(output_file),
                            format='FFMPEG',
                            mode='I',
                            fps=fps,
                            codec='libx264',
                            pixelformat='yuv420p',
                            output_params=['-crf', '18']
                        )

                        for frame in frames:
                            # 转换BGR到RGB
                            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                            writer.append_data(frame_rgb)

                        writer.close()
                        print(f"✅ 视频导出完成: {output_file}")
                        relative_path = output_file.relative_to(self.core_space_dir)
                        self.gui_controls['export_status'].value = f"完成! {relative_path}"

                    except ImportError:
                        # 使用OpenCV写入
                        height, width = frames[0].shape[:2]
                        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                        writer = cv2.VideoWriter(str(output_file), fourcc, fps, (width, height))
                        for frame in frames:
                            writer.write(frame)
                        writer.release()
                        print(f"✅ 视频导出完成: {output_file}")
                        relative_path = output_file.relative_to(self.core_space_dir)
                        self.gui_controls['export_status'].value = f"完成! {relative_path}"

                except Exception as e:
                    print(f"❌ Selenium截图失败: {e}")
                    import traceback
                    traceback.print_exc()
                    # 回退到渲染方法
                    use_selenium = False

            if not use_selenium:
                # 回退到原来的渲染方法
                print("   使用PyRender离线渲染...")
                self._export_video_thread()
                return

        except Exception as e:
            print(f"❌ 导出视频失败: {e}")
            import traceback
            traceback.print_exc()
            self.gui_controls['export_status'].value = f"错误: {str(e)}"

        finally:
            self.is_exporting = False

    def _render_frame_offline(self, frame_idx: int, resolution: int,
                              camera_pos: np.ndarray, camera_wxyz: np.ndarray) -> Optional[np.ndarray]:
        """离线渲染一帧"""
        # 尝试导入pyrender
        try:
            import pyrender
            import trimesh
        except ImportError:
            if frame_idx == 0:
                print("⚠️ pyrender未安装，使用简化渲染...")
                print("   提示: 安装 pyrender 以获得完整3D渲染")
                print("   pip install pyrender trimesh")
            return self._render_frame_simple(frame_idx, resolution)

        # 设置PyRender使用离屏渲染（EGL或OSMesa）
        # 优先尝试EGL，如果失败则尝试OSMesa
        for platform in ['egl', 'osmesa']:
            try:
                os.environ['PYOPENGL_PLATFORM'] = platform

                # 创建场景 - 设置深蓝色背景（与viser一致）
                scene = pyrender.Scene(
                    ambient_light=[0.3, 0.3, 0.3],
                    bg_color=[13/255, 13/255, 38/255, 1.0]  # 深蓝色背景
                )

                # 获取GUI参数
                show_generated = self.gui_controls['show_generated'].value
                show_gt = self.gui_controls['show_gt'].value
                generated_color = np.array(self.gui_controls['generated_color'].value) / 255.0
                gt_color = np.array(self.gui_controls['gt_color'].value) / 255.0
                mesh_resolution = int(self.gui_controls['mesh_resolution'].value)

                mesh_count = 0

                # 添加生成的超二次曲面
                # 重要：需要应用场景归一化，使物体坐标与相机坐标在同一空间
                if show_generated:
                    predictions = self._extract_predictions(frame_idx)
                    if predictions is not None:
                        for obj_idx, obj_params in enumerate(predictions):
                            if obj_params[0] > 0.5:
                                # 复制参数并应用场景归一化到平移部分
                                obj_params_normalized = obj_params.copy()
                                # 归一化平移: (translation - scene_center) * scene_scale
                                translation = obj_params[6:9]
                                translation_normalized = (translation - self.scene_center) * self.scene_scale
                                obj_params_normalized[6:9] = translation_normalized
                                # 归一化缩放: scale * scene_scale
                                obj_params_normalized[3:6] = obj_params[3:6] * self.scene_scale

                                vertices, faces = self.generate_superquadric_mesh(
                                    obj_params_normalized, num_samples=mesh_resolution
                                )

                                if frame_idx == 0 and obj_idx == 0:
                                    print(f"   物体原始位置: {translation}")
                                    print(f"   物体归一化位置: {translation_normalized}")
                                    print(f"   场景中心: {self.scene_center}, 缩放: {self.scene_scale}")

                                mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
                                # 为每个顶点设置颜色 (N, 4) - RGBA
                                num_verts = len(vertices)
                                vertex_colors = np.zeros((num_verts, 4), dtype=np.uint8)
                                vertex_colors[:, :3] = (generated_color * 255).astype(np.uint8)  # RGB
                                vertex_colors[:, 3] = 255  # 完全不透明
                                mesh.visual.vertex_colors = vertex_colors

                                # 创建PyRender材质
                                material = pyrender.MetallicRoughnessMaterial(
                                    baseColorFactor=list(generated_color) + [1.0],
                                    metallicFactor=0.3,
                                    roughnessFactor=0.7
                                )
                                mesh_obj = pyrender.Mesh.from_trimesh(mesh, material=material)
                                scene.add(mesh_obj)
                                mesh_count += 1

                # 添加GT超二次曲面
                if show_gt:
                    targets = self._extract_targets(frame_idx)
                    if targets is not None:
                        for obj_idx, obj_params in enumerate(targets):
                            if obj_params[0] > 0.5:
                                # 复制参数并应用场景归一化
                                obj_params_normalized = obj_params.copy()
                                translation = obj_params[6:9]
                                translation_normalized = (translation - self.scene_center) * self.scene_scale
                                obj_params_normalized[6:9] = translation_normalized
                                obj_params_normalized[3:6] = obj_params[3:6] * self.scene_scale

                                vertices, faces = self.generate_superquadric_mesh(
                                    obj_params_normalized, num_samples=mesh_resolution
                                )
                                mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
                                # 为每个顶点设置颜色 (N, 4) - RGBA
                                num_verts = len(vertices)
                                vertex_colors = np.zeros((num_verts, 4), dtype=np.uint8)
                                vertex_colors[:, :3] = (gt_color * 255).astype(np.uint8)  # RGB
                                vertex_colors[:, 3] = 255  # 完全不透明
                                mesh.visual.vertex_colors = vertex_colors

                                # 创建PyRender材质
                                material = pyrender.MetallicRoughnessMaterial(
                                    baseColorFactor=list(gt_color) + [0.5],
                                    metallicFactor=0.3,
                                    roughnessFactor=0.7
                                )
                                mesh_obj = pyrender.Mesh.from_trimesh(mesh, material=material)
                                scene.add(mesh_obj)
                                mesh_count += 1

                if frame_idx == 0:
                    print(f"   场景中添加了 {mesh_count} 个mesh")

                # 设置相机
                # Viser使用的是wxyz四元数，需要转换为PyRender的变换矩阵
                from scipy.spatial.transform import Rotation as R

                # wxyz -> xyzw for scipy
                rot = R.from_quat([camera_wxyz[1], camera_wxyz[2], camera_wxyz[3], camera_wxyz[0]])
                rot_matrix = rot.as_matrix()

                # PyRender使用OpenGL坐标系
                # 构建相机变换矩阵
                camera_pose = np.eye(4)
                camera_pose[:3, :3] = rot_matrix
                camera_pose[:3, 3] = camera_pos

                if frame_idx == 0:
                    print(f"   相机位置: {camera_pos}")
                    print(f"   相机旋转矩阵:\n{rot_matrix}")

                # 创建透视相机
                camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.0)
                scene.add(camera, pose=camera_pose)

                # 添加多个光源以确保场景被充分照亮
                # 主光源跟随相机
                light1 = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=2.0)
                scene.add(light1, pose=camera_pose)

                # 额外的环境光源
                light2 = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1.0)
                light_pose = np.eye(4)
                light_pose[:3, 3] = [10, 10, 10]
                scene.add(light2, pose=light_pose)

                # 渲染
                renderer = pyrender.OffscreenRenderer(resolution, resolution)
                color, depth = renderer.render(scene)
                renderer.delete()

                # 首次成功时打印使用的平台和渲染统计
                if frame_idx == 0:
                    print(f"   ✅ 使用 {platform.upper()} 进行离线渲染")
                    print(f"   渲染输出范围: [{color.min()}, {color.max()}]")
                    print(f"   深度范围: [{depth.min()}, {depth.max()}]")

                return color

            except Exception as e:
                if platform == 'osmesa':
                    # 两种方式都失败了
                    if frame_idx == 0:
                        print(f"❌ PyRender渲染失败 (EGL和OSMesa都不可用): {e}")
                        print("   使用简化渲染模式...")
                    return self._render_frame_simple(frame_idx, resolution)
                # EGL失败，继续尝试OSMesa
                continue

        # 不应该到达这里，但以防万一
        return self._render_frame_simple(frame_idx, resolution)

    def _render_frame_simple(self, frame_idx: int, resolution: int) -> np.ndarray:
        """简化渲染（纯色背景 + 文字提示）"""
        # 创建空白图像
        image = np.full((resolution, resolution, 3), [13, 13, 38], dtype=np.uint8)

        # 添加文字
        text = f"Frame {frame_idx + 1}/{self.original_frame_count}"
        font = cv2.FONT_HERSHEY_SIMPLEX
        text_size = cv2.getTextSize(text, font, 1, 2)[0]
        text_x = (resolution - text_size[0]) // 2
        text_y = (resolution + text_size[1]) // 2

        cv2.putText(image, text, (text_x, text_y), font, 1, (255, 255, 255), 2)

        # 添加提示信息
        hint = "Install pyrender for full rendering"
        hint_size = cv2.getTextSize(hint, font, 0.5, 1)[0]
        hint_x = (resolution - hint_size[0]) // 2
        hint_y = text_y + 40

        cv2.putText(image, hint, (hint_x, hint_y), font, 0.5, (150, 150, 150), 1)

        return image

    def run(self, auto_open_browser: bool = True):
        """运行可视化器"""
        print("\n" + "="*60)
        print("🎨 WaveGen 训练可视化器")
        print("="*60)
        print(f"📁 监控目录: {self.core_space_dir}")
        print(f"🌐 Web界面: http://localhost:{self.port}")
        print("="*60)
        print("\n💡 提示:")
        print("   - 如果页面空白一直加载，请刷新浏览器 (Ctrl+Shift+R)")
        print("   - 建议使用 Chrome 或 Firefox 浏览器")
        print("\n按 Ctrl+C 退出\n")

        # 自动打开浏览器
        if auto_open_browser:
            url = f"http://localhost:{self.port}"
            print(f"🌐 正在打开浏览器: {url}")
            try:
                webbrowser.open(url)
            except Exception as e:
                print(f"⚠️  无法自动打开浏览器: {e}")
                print(f"   请手动访问: {url}")

        try:
            while True:
                time.sleep(0.1)
        except KeyboardInterrupt:
            print("\n👋 再见！")
            print("正在关闭服务器...")
            # 清理资源
            try:
                for mesh in self.mesh_handles_pool.values():
                    mesh.remove()
            except:
                pass


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description="WaveGen训练结果可视化工具")
    parser.add_argument(
        '--core-space',
        type=str,
        default='core_space',
        help='core_space目录路径（默认: ./core_space）'
    )
    parser.add_argument(
        '--port',
        type=int,
        default=8080,
        help='Viser服务器端口（默认: 8080，如果被占用会自动尝试下一个端口）'
    )
    parser.add_argument(
        '--no-browser',
        action='store_true',
        help='不自动打开浏览器'
    )

    args = parser.parse_args()

    # 创建并运行可视化器
    visualizer = TrainingVisualizer(core_space_dir=args.core_space, port=args.port)
    visualizer.run(auto_open_browser=not args.no_browser)


if __name__ == "__main__":
    main()