nano_WaveGen/utils/save_generation_results.py

"""
保存生成结果的工具函数
用于WaveGen v33 - 使用超二次元函数的版本
"""
import os
import numpy as np
import torch
from pathlib import Path
from typing import Dict, List, Optional
import json
from datetime import datetime
import shutil


def save_generation_results(
    predictions: List[Dict],
    targets: Dict[str, torch.Tensor],
    texts: List[str],
    step: int,
    output_dir: str = "outputs",
    save_config: Dict = None,
    metadata: Dict = None,
    batch_data: Dict = None,  # 新增：完整的批次数据
    data_root: str = None,     # 新增：原始数据根目录
    data_split: str = "validation"  # 新增：数据集split（train/validation）
):
    """
    保存生成结果用于可视化和分析（增强版）
    
    Args:
        predictions: 模型预测结果列表，每个元素包含 'objects', 'world', 'physics'
        targets: 真实目标数据
        texts: 输入文本描述
        step: 当前训练步数
        output_dir: 输出目录
        save_config: 保存配置
        metadata: 额外的元数据（如序列名称、相机参数等）
        batch_data: 完整的批次数据，用于获取更多原始信息
        data_root: 原始数据根目录，用于复制原始文件
        data_split: 数据集split（'train' 或 'validation'），用于确定原始数据的正确位置
    """
    # 创建时间戳目录
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    save_path = Path(output_dir) / f"{timestamp}_step{step}_text2wave"
    save_path.mkdir(parents=True, exist_ok=True)
    
    # 保存配置
    if save_config is not None:
        with open(save_path / "save_config.json", 'w') as f:
            json.dump(save_config, f, indent=2)
    
    batch_size = len(texts)
    num_frames = len(predictions)

    print(f"\n{'='*60}")
    print(f"💾 保存生成结果: step {step}")
    print(f"📁 保存路径: {save_path}")
    print(f"📊 样本数: {batch_size}, 帧数: {num_frames}")
    print(f"{'='*60}\n")

    for i in range(batch_size):
        try:
            sample_dir = save_path / f"sample_{i}"
            sample_dir.mkdir(exist_ok=True)
            print(f"正在保存 sample_{i}... ", end='', flush=True)

            # 1. 保存文本描述和元信息
            with open(sample_dir / "info.txt", 'w') as f:
                f.write(f"Text: {texts[i]}\n")
                f.write(f"Generated at step: {step}\n")
                f.write(f"Number of frames: {num_frames}\n")

                if metadata and 'sequence_names' in metadata and metadata['sequence_names'] is not None:
                    f.write(f"Sequence: {metadata['sequence_names'][i]}\n")

                f.write("\n--- Model Output Summary ---\n")
                f.write(f"Max objects: {predictions[0]['objects'].shape[1]}\n")
                f.write("Object parameters: 15 (exists + shape[2] + scale[3] + translation[3] + rotation[3] + velocity[3])\n")
                f.write(f"World parameters: 8 (camera_pos[3] + camera_quat[4] + scene_scale[1])\n")
                f.write(f"Physics parameters: 3 (mass + friction + restitution)\n")

            # 2. 保存生成的超二次元函数参数（改进格式）
            frame_predictions = []
            physics_per_frame = []  # 保存每一帧的physics预测，便于查看全序列
            for f_idx, pred in enumerate(predictions):
                # predictions列表中每个元素包含当前帧、完整batch的数据，这里先取出当前样本
                objects_batch = pred['objects']
                world_batch = pred['world']
                physics_batch = pred.get('physics')

                objects_params = objects_batch[i].cpu().numpy() if hasattr(objects_batch, 'cpu') else objects_batch[i]  # [max_objects, 15]
                world_params = world_batch[i].cpu().numpy() if hasattr(world_batch, 'cpu') else world_batch[i]  # [8]
                if physics_batch is not None:
                    physics_params = physics_batch[i].cpu().numpy() if hasattr(physics_batch, 'cpu') else physics_batch[i]  # [max_objects, 3]
                else:
                    physics_params = np.zeros((objects_params.shape[0], 3), dtype=np.float32)

                # 保存当前帧的 physics（每帧各物体的mass/friction/restitution）
                physics_per_frame.append([
                    {
                        'mass': float(phys_params[0]),
                        'friction': float(phys_params[1]),
                        'restitution': float(phys_params[2]),
                    }
                    for phys_params in physics_params
                ])

                # 将物体参数转换为更易读的格式
                superquadrics = []
                for obj_idx in range(objects_params.shape[0]):
                    obj_params = objects_params[obj_idx]
                    phys_params = physics_params[obj_idx]

                    superquadric = {
                        'exists': bool(obj_params[0] > 0.5),  # exists flag
                        'shape': obj_params[1:3],              # epsilon1, epsilon2
                        'scale': obj_params[3:6],              # a, b, c
                        'translation': obj_params[6:9],        # x, y, z
                        'rotation': obj_params[9:12],          # euler angles
                        # 预测没有 inlier_ratio，填充0以保持键一致
                        'inlier_ratio': 0.0,
                        'velocity': obj_params[12:15],         # vx, vy, vz
                        'mass': phys_params[0],
                        'friction': phys_params[1],
                        'restitution': phys_params[2],
                    }
                    superquadrics.append(superquadric)

                # 将世界参数转换为更易读的格式
                world_info = {
                    'camera_position': world_params[0:3],      # x, y, z
                    'camera_quaternion': world_params[3:7],    # w, x, y, z
                    'scene_scale': float(world_params[7]),     # scale
                    # 预测没有scene_center，填零保持字段一致，方便下游读取
                    'scene_center': np.zeros(3, dtype=np.float32),
                }

                frame_data = {
                    'frame_idx': f_idx,
                    'superquadrics': superquadrics,
                    'world_info': world_info,
                }
                frame_predictions.append(frame_data)

            np.savez(sample_dir / "predictions.npz",
                     text=texts[i],
                     frames=frame_predictions,
                     num_frames=num_frames,
                     # 保存全序列physics；未预测则写None
                     physics=physics_per_frame if physics_per_frame else None,
                     sequence_name=metadata['sequence_names'][i] if (metadata and 'sequence_names' in metadata and metadata['sequence_names'] is not None) else "unknown",
                     description="Predicted superquadric parameters for each frame")

            # 3. 保存真实目标数据（如果有）- 改进格式
            if targets is not None:
                # targets中的数据已经是完整批次，需要索引[i]获取当前样本
                target_objects = targets['objects'][i].cpu().numpy() if hasattr(targets['objects'], 'cpu') else targets['objects'][i]   # [num_frames, max_objects, 16]
                target_world = targets['world'][i].cpu().numpy() if hasattr(targets['world'], 'cpu') else targets['world'][i]       # [num_frames, 11]

                if 'physics' in targets and targets['physics'] is not None:
                    target_physics = targets['physics'][i].cpu().numpy() if hasattr(targets['physics'], 'cpu') else targets['physics'][i]
                else:
                    target_physics = None

                # 生成顶层 physics，与原始 Full_Sample_Data_for_Learning_Target 一致
                target_physics_top = None
                if target_physics is not None:
                    target_physics_top = [
                        {
                            'mass': float(p[0]),
                            'friction': float(p[1]),
                            'restitution': float(p[2]),
                        }
                        for p in target_physics
                    ]

                # 将目标数据转换为更易读的格式
                target_frames = []
                for f_idx in range(target_objects.shape[0]):
                    frame_objects = target_objects[f_idx]  # [max_objects, 16]
                    frame_world = target_world[f_idx]      # [11]

                    # 转换物体参数
                    superquadrics = []
                    for obj_idx in range(frame_objects.shape[0]):
                        obj_params = frame_objects[obj_idx]
                        phys_params = target_physics[obj_idx] if target_physics is not None else np.zeros(3)

                        superquadric = {
                            'exists': bool(obj_params[0] > 0.5),   # exists flag
                            'shape': obj_params[1:3],               # epsilon1, epsilon2
                            'scale': obj_params[3:6],               # a, b, c
                            'translation': obj_params[6:9],         # x, y, z
                            'rotation': obj_params[9:12],           # euler angles
                            'inlier_ratio': float(obj_params[12]),  # GT specific: inlier ratio
                            'velocity': obj_params[13:16],          # vx, vy, vz
                            'mass': phys_params[0],
                            'friction': phys_params[1],
                            'restitution': phys_params[2],
                        }
                        superquadrics.append(superquadric)

                    # 转换世界参数
                    world_info = {
                        'camera_position': frame_world[0:3],       # x, y, z
                        'camera_quaternion': frame_world[3:7],     # w, x, y, z
                        'scene_scale': float(frame_world[7]),      # scale
                        'scene_center': frame_world[8:11],         # center x, y, z
                    }

                    frame_data = {
                        'frame_idx': f_idx,
                        'superquadrics': superquadrics,
                        'world_info': world_info,
                    }
                    target_frames.append(frame_data)

                # 保存改进格式的 targets.npz
                np.savez(sample_dir / "targets.npz",
                         text=texts[i],
                         frames=target_frames,
                         num_frames=num_frames,
                         physics=target_physics_top if target_physics_top is not None else None,
                         sequence_name=metadata['sequence_names'][i] if (metadata and 'sequence_names' in metadata and metadata['sequence_names'] is not None) else "unknown",
                         description="Ground truth superquadric parameters for each frame")

                # 为了兼容性，也保存原始格式（用于误差计算）
                target_data_legacy = {
                    'objects': target_objects,
                    'world': target_world,
                    'physics': target_physics,
                }

                # 计算并保存误差统计（使用旧格式）
                save_error_statistics(frame_predictions, target_data_legacy, sample_dir)

            # 4. 保存相机参数（如果有）
            if metadata and 'camera_data' in metadata:
                camera_data = metadata['camera_data'][i]
                np.savez(sample_dir / "camera_params.npz",
                         **camera_data)

            # 5. 保存原始数据（新增功能，不再依赖camera_data）
            if batch_data is not None and data_root is not None:
                save_original_data(sample_dir, i, batch_data, metadata, data_root, data_split)

            print("✅")  # 完成标记

        except Exception as e:
            print(f"❌ 错误: {e}")
            import traceback
            traceback.print_exc()
            continue  # 继续保存其他样本

        # 注意：已移除save_visualization_script，因为不需要单独的可视化脚本
    
    # 保存整体统计信息
    save_batch_statistics(predictions, targets, save_path)

    print(f"\n{'='*60}")
    print(f"✅ 保存完成!")
    print(f"📁 保存路径: {save_path}")
    print(f"{'='*60}\n")

    return save_path


def save_error_statistics(predictions: List[Dict], targets: Dict, save_dir: Path):
    """计算并保存预测误差统计

    Args:
        predictions: 新格式的帧列表 (包含 superquadrics 和 world_info)
        targets: 旧格式的目标数据 (包含 objects, world 数组)
    """
    stats = {}

    # 将新格式的 predictions 转换回数组格式进行误差计算
    object_errors = []
    world_errors = []

    for frame in predictions:
        frame_idx = frame['frame_idx']

        # 从新格式重建物体数组
        superquadrics = frame['superquadrics']
        pred_objects = []
        for sq in superquadrics:
            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']
            obj_params[3:6] = sq['scale']
            obj_params[6:9] = sq['translation']
            obj_params[9:12] = sq['rotation']
            obj_params[12:15] = sq['velocity']
            pred_objects.append(obj_params)
        pred_obj = np.array(pred_objects)

        # 获取目标物体数据
        target_obj_full = targets['objects'][frame_idx]
        target_obj = target_obj_full[:, :pred_obj.shape[1]]  # 对齐模型预测的维度

        # 只计算存在的物体
        exists_mask = target_obj[:, 0] > 0.5
        if exists_mask.any():
            error = np.mean(np.abs(pred_obj[exists_mask] - target_obj[exists_mask]))
            object_errors.append(error)

        # 从新格式重建世界参数数组
        world_info = frame['world_info']
        pred_world = np.concatenate([
            world_info['camera_position'],
            world_info['camera_quaternion'],
            [world_info['scene_scale']]
        ])

        # 获取目标世界数据
        target_world = targets['world'][frame_idx][:8]
        error = np.mean(np.abs(pred_world - target_world))
        world_errors.append(error)

    stats['object_mae'] = float(np.mean(object_errors)) if object_errors else 0.0
    stats['world_mae'] = float(np.mean(world_errors))

    # 保存统计信息
    with open(save_dir / "error_statistics.json", 'w') as f:
        json.dump(stats, f, indent=2)


def save_batch_statistics(predictions: List[Dict], targets: Dict, save_dir: Path):
    """保存整批数据的统计信息"""
    batch_size = predictions[0]['objects'].shape[0]
    
    stats = {
        'batch_size': batch_size,
        'num_frames': len(predictions),
        'timestamp': datetime.now().isoformat(),
    }
    
    # 统计每帧实际存在的物体数量
    if targets is not None:
        objects_per_frame = []
        for f_idx in range(len(predictions)):
            frame_objects = []
            for b_idx in range(batch_size):
                exists = targets['objects'][b_idx, f_idx, :, 0] > 0.5
                frame_objects.append(int(exists.sum()))
            objects_per_frame.append({
                'frame': f_idx,
                'mean_objects': float(np.mean(frame_objects)),
                'max_objects': int(max(frame_objects)),
                'min_objects': int(min(frame_objects)),
            })
        stats['objects_per_frame'] = objects_per_frame
    
    with open(save_dir / "batch_statistics.json", 'w') as f:
        json.dump(stats, f, indent=2)


def save_visualization_script(save_dir: Path):
    """保存用于可视化超二次元函数的Python脚本"""
    script = '''#!/usr/bin/env python3
"""
可视化生成的超二次元函数参数
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def superquadric_surface(a, b, c, e1, e2, n=50):
    """生成超二次元函数表面点"""
    eta = np.linspace(-np.pi/2, np.pi/2, n)
    omega = np.linspace(-np.pi, np.pi, n)
    eta, omega = np.meshgrid(eta, omega)
    
    x = a * np.sign(np.cos(eta)) * np.abs(np.cos(eta))**e1 * np.sign(np.cos(omega)) * np.abs(np.cos(omega))**e2
    y = b * np.sign(np.cos(eta)) * np.abs(np.cos(eta))**e1 * np.sign(np.sin(omega)) * np.abs(np.sin(omega))**e2
    z = c * np.sign(np.sin(eta)) * np.abs(np.sin(eta))**e1
    
    return x, y, z

# 加载预测数据
data = np.load('predictions.npz', allow_pickle=True)
frames = data['frames']

# 可视化第一帧
frame = frames[0]
objects = frame['objects']  # [max_objects, 12]

fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')

# 绘制每个存在的物体
for obj_idx, obj_params in enumerate(objects):
    if obj_params[0] > 0.5:  # 物体存在
        # 提取参数
        shape_params = obj_params[1:3]  # e1, e2
        scale = obj_params[3:6]  # a, b, c
        translation = obj_params[6:9]
        rotation = obj_params[9:12]  # 简化处理，暂不应用旋转
        
        # 生成表面
        x, y, z = superquadric_surface(scale[0], scale[1], scale[2], 
                                      shape_params[0], shape_params[1])
        
        # 应用平移
        x += translation[0]
        y += translation[1]
        z += translation[2]
        
        # 绘制
        ax.plot_surface(x, y, z, alpha=0.7, label=f'Object {obj_idx}')

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_title('Generated Superquadric Objects')
plt.savefig('visualization.png')
plt.show()
'''
    
    with open(save_dir / "visualize.py", 'w') as f:
        f.write(script)
    
    # 使脚本可执行
    os.chmod(save_dir / "visualize.py", 0o755)


def save_original_data(sample_dir: Path, sample_idx: int, batch_data: Dict, metadata: Dict, data_root: str, data_split: str = "validation"):
    """
    保存原始数据文件，包括RGB图像、深度图、分割图、点云等
    
    Args:
        sample_dir: 当前样本保存目录
        sample_idx: 批次中的样本索引
        batch_data: 批次数据
        metadata: 元数据
        data_root: MOVi数据集根目录
        data_split: 数据集split（'train' 或 'validation'）
    """
    try:
        # 获取原始序列名称
        sequence_name = None
        if metadata and 'sequence_names' in metadata and metadata['sequence_names'] is not None:
            sequence_name = metadata['sequence_names'][sample_idx]
        
        if not sequence_name:
            print(f"Warning: No sequence name found for sample {sample_idx}")
            # 创建一个说明文件，解释为什么没有原始数据
            no_original_data_file = sample_dir / "no_original_data.txt"
            with open(no_original_data_file, 'w') as f:
                f.write("原始数据未保存，因为无法获取序列名称。\n")
                f.write("这可能是因为数据加载器没有返回sequence_names字段。\n")
                f.write(f"Sample index: {sample_idx}\n")
                f.write(f"Generated at step: {datetime.now().isoformat()}\n")
            return
            
        # 构建原始数据路径
        data_root_path = Path(data_root)
        
        # 直接使用指定的data_split来定位原始数据
        original_sample_dir = data_root_path / data_split / sequence_name
        
        if not original_sample_dir.exists():
            print(f"Warning: Could not find original data for {sequence_name} in {data_split} split")
            # 创建一个说明文件
            error_file = sample_dir / "original_data_not_found.txt"
            with open(error_file, 'w') as f:
                f.write(f"原始数据未找到\n")
                f.write(f"查找路径: {original_sample_dir}\n")
                f.write(f"数据集: {data_split}\n")
                f.write(f"序列名: {sequence_name}\n")
                f.write(f"时间: {datetime.now().isoformat()}\n")
            return
            
        print(f"Copying original data from {original_sample_dir}")
        
        # 创建原始数据子目录
        original_data_dir = sample_dir / "original_data"
        original_data_dir.mkdir(exist_ok=True)
        
        # 1. 复制RGB图像（所有帧）
        rgb_dir = original_data_dir / "rgb"
        rgb_dir.mkdir(exist_ok=True)
        original_rgb_dir = original_sample_dir / "rgb"
        if original_rgb_dir.exists():
            for rgb_file in sorted(original_rgb_dir.glob("frame_*.png")):
                shutil.copy2(rgb_file, rgb_dir / rgb_file.name)
        
        # 2. 复制深度图（优先复制合并的npz，否则复制单独的npy）
        depth_dir = original_data_dir / "depth"
        depth_dir.mkdir(exist_ok=True)
        original_depth_dir = original_sample_dir / "depth"
        if original_depth_dir.exists():
            # 检查是否有合并的npz文件
            merged_depth = original_depth_dir / "depth_merge.npz"
            if merged_depth.exists():
                shutil.copy2(merged_depth, depth_dir / "depth_merge.npz")
            else:
                # 复制单独的npy文件
                for depth_file in sorted(original_depth_dir.glob("frame_*.npy")):
                    shutil.copy2(depth_file, depth_dir / depth_file.name)

        # 3. 复制分割图（优先复制合并的npz，否则复制单独的npy）
        seg_dir = original_data_dir / "segmentation"
        seg_dir.mkdir(exist_ok=True)
        original_seg_dir = original_sample_dir / "segmentation"
        if original_seg_dir.exists():
            # 检查是否有合并的npz文件
            merged_seg = original_seg_dir / "segmentation_merge.npz"
            if merged_seg.exists():
                shutil.copy2(merged_seg, seg_dir / "segmentation_merge.npz")
            else:
                # 复制单独的npy文件
                for seg_file in sorted(original_seg_dir.glob("frame_*.npy")):
                    shutil.copy2(seg_file, seg_dir / seg_file.name)

        # 4. 复制法线图（优先复制合并的npz，否则复制单独的npy）
        normal_dir = original_data_dir / "normal"
        normal_dir.mkdir(exist_ok=True)
        original_normal_dir = original_sample_dir / "normal"
        if original_normal_dir.exists():
            # 检查是否有合并的npz文件
            merged_normal = original_normal_dir / "normal_merge.npz"
            if merged_normal.exists():
                shutil.copy2(merged_normal, normal_dir / "normal_merge.npz")
            else:
                # 复制单独的npy文件
                for normal_file in sorted(original_normal_dir.glob("frame_*.npy")):
                    shutil.copy2(normal_file, normal_dir / normal_file.name)
                
        # 5. 复制相机轨迹
        camera_traj_file = original_sample_dir / "camera_trajectory.npz"
        if camera_traj_file.exists():
            shutil.copy2(camera_traj_file, original_data_dir / "camera_trajectory.npz")
            
        # 6. 复制元数据JSON
        metadata_file = original_sample_dir / "metadata.json"
        if metadata_file.exists():
            shutil.copy2(metadata_file, original_data_dir / "metadata.json")

        # 7. 复制完整的训练目标数据缓存文件（Full_Sample_Data_for_Learning_Target.npz）
        full_cache_file = original_sample_dir / "Full_Sample_Data_for_Learning_Target.npz"
        if full_cache_file.exists():
            shutil.copy2(full_cache_file, original_data_dir / "Full_Sample_Data_for_Learning_Target.npz")

        # 8. 复制其他可能的合并文件（object_coordinates, point_clouds等）
        for folder_name in ['object_coordinates', 'point_clouds']:
            folder_dir = original_data_dir / folder_name
            original_folder_dir = original_sample_dir / folder_name
            if original_folder_dir.exists():
                folder_dir.mkdir(exist_ok=True)
                # 检查是否有合并的npz文件
                merged_file = original_folder_dir / f"{folder_name}_merge.npz"
                if merged_file.exists():
                    shutil.copy2(merged_file, folder_dir / f"{folder_name}_merge.npz")
                else:
                    # 复制单独的npy文件
                    for npy_file in sorted(original_folder_dir.glob("frame_*.npy")):
                        shutil.copy2(npy_file, folder_dir / npy_file.name)

        # 9. 如果有预处理的点云数据（在batch_data中），也保存
        if 'point_clouds' in batch_data:
            pc_data = batch_data['point_clouds'][sample_idx]
            np.savez_compressed(original_data_dir / "point_clouds.npz", **pc_data)

        # 10. 保存场景归一化参数
        if 'scene_normalization' in batch_data:
            norm_params = batch_data['scene_normalization'][sample_idx]
            with open(original_data_dir / "scene_normalization.json", 'w') as f:
                json.dump({
                    'scene_center': norm_params['center'].tolist() if hasattr(norm_params['center'], 'tolist') else norm_params['center'],
                    'scene_scale': float(norm_params['scale']) if 'scale' in norm_params else 1.0,
                    'scene_extent': float(norm_params['extent']) if 'extent' in norm_params else 1.0
                }, f, indent=2)

        # 11. 创建文件清单
        with open(original_data_dir / "file_manifest.txt", 'w') as f:
            f.write(f"Original sequence: {sequence_name}\n")
            f.write(f"Data split: {data_split}\n")
            f.write(f"Original path: {original_sample_dir}\n")
            f.write(f"Copied at: {datetime.now().isoformat()}\n\n")
            f.write("Files included:\n")
            for item in sorted(original_data_dir.rglob("*")):
                if item.is_file() and item.name != "file_manifest.txt":
                    f.write(f"- {item.relative_to(original_data_dir)}\n")
                    
    except Exception as e:
        print(f"Error saving original data for sample {sample_idx}: {e}")
        import traceback
        traceback.print_exc()