scripts/encode_spatialvid.py

import os
import torch
import lightning as pl
from PIL import Image
from diffsynth import WanVideoReCamMasterPipeline, ModelManager
import json
import imageio
from torchvision.transforms import v2
from einops import rearrange
import argparse
import numpy as np
import pdb
from tqdm import tqdm
import pandas as pd

os.environ["TOKENIZERS_PARALLELISM"] = "false"

from scipy.spatial.transform import Slerp
from scipy.spatial.transform import Rotation as R

def interpolate_camera_poses(original_frames, original_poses, target_frames):
    """
    对相机姿态进行插值，生成目标帧对应的姿态参数
    
    参数:
    original_frames: 原始帧索引列表，如[0,6,12,...]
    original_poses: 原始姿态数组，形状为(n,7)，每行[tx, ty, tz, qx, qy, qz, qw]
    target_frames: 目标帧索引列表，如[0,4,8,12,...]
    
    返回:
    target_poses: 插值后的姿态数组，形状为(m,7),m为目标帧数量
    """
    # 确保输入有效
    print('original_frames:',len(original_frames))
    print('original_poses:',len(original_poses))
    if len(original_frames) != len(original_poses):
        raise ValueError("原始帧数量与姿态数量不匹配")
    
    if original_poses.shape[1] != 7:
        raise ValueError(f"原始姿态应为(n,7)格式，实际为{original_poses.shape}")
    
    target_poses = []
    
    # 提取旋转部分并转换为Rotation对象
    rotations = R.from_quat(original_poses[:, 3:7])  # 提取四元数部分
    
    for t in target_frames:
        # 找到t前后的原始帧索引
        idx = np.searchsorted(original_frames, t, side='left')
        
        # 处理边界情况
        if idx == 0:
            # 使用第一个姿态
            target_poses.append(original_poses[0])
            continue
        if idx >= len(original_frames):
            # 使用最后一个姿态
            target_poses.append(original_poses[-1])
            continue
            
        # 获取前后帧的信息
        t_prev, t_next = original_frames[idx-1], original_frames[idx]
        pose_prev, pose_next = original_poses[idx-1], original_poses[idx]
        
        # 计算插值权重
        alpha = (t - t_prev) / (t_next - t_prev)
        
        # 1. 平移向量的线性插值
        translation_prev = pose_prev[:3]
        translation_next = pose_next[:3]
        interpolated_translation = translation_prev + alpha * (translation_next - translation_prev)
        
        # 2. 旋转四元数的球面线性插值(SLERP)
        # 创建Slerp对象
        slerp = Slerp([t_prev, t_next], rotations[idx-1:idx+1])
        interpolated_rotation = slerp(t)
        
        # 组合平移和旋转
        interpolated_pose = np.concatenate([
            interpolated_translation,
            interpolated_rotation.as_quat()  # 转换回四元数
        ])
        
        target_poses.append(interpolated_pose)
    
    return np.array(target_poses)


class VideoEncoder(pl.LightningModule):
    def __init__(self, text_encoder_path, vae_path, tiled=True, tile_size=(34, 34), tile_stride=(18, 16)):
        super().__init__()
        model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
        model_manager.load_models([text_encoder_path, vae_path])
        self.pipe = WanVideoReCamMasterPipeline.from_model_manager(model_manager)
        self.tiler_kwargs = {"tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride}
        
        self.frame_process = v2.Compose([
            # v2.CenterCrop(size=(900, 1600)),
            # v2.Resize(size=(900, 1600), antialias=True),
            v2.ToTensor(),
            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        ])
    
    def crop_and_resize(self, image):
        width, height = image.size
        # print(width,height)
        width_ori, height_ori_ = 832 , 480
        image = v2.functional.resize(
            image,
            (round(height_ori_), round(width_ori)),
            interpolation=v2.InterpolationMode.BILINEAR
        )
        return image

    def load_video_frames(self, video_path):
        """加载完整视频"""
        reader = imageio.get_reader(video_path)
        frames = []
        
        for frame_data in reader:
            frame = Image.fromarray(frame_data)
            frame = self.crop_and_resize(frame)
            frame = self.frame_process(frame)
            frames.append(frame)
        
        reader.close()
        
        if len(frames) == 0:
            return None
            
        frames = torch.stack(frames, dim=0)
        frames = rearrange(frames, "T C H W -> C T H W")
        return frames

def encode_scenes(scenes_path, text_encoder_path, vae_path,output_dir):
    """编码所有场景的视频"""

    encoder = VideoEncoder(text_encoder_path, vae_path)
    encoder = encoder.cuda()
    encoder.pipe.device = "cuda"
    
    processed_count = 0

    processed_chunk_count = 0

    prompt_emb = 0

    metadata = pd.read_csv('/share_zhuyixuan05/public_datasets/SpatialVID-HQ/data/train/SpatialVID_HQ_metadata.csv')


    os.makedirs(output_dir,exist_ok=True)
    chunk_size = 300
    required_keys = ["latents", "cam_emb", "prompt_emb"]

    for i, scene_name in enumerate(os.listdir(scenes_path)):
        # print('index-----:',type(i))
        if i < 3 :#or i >=2000:
        #     # print('index-----:',i)
            continue
            # print('index:',i)
        print('group:',i)
        scene_dir = os.path.join(scenes_path, scene_name)
        
        # save_dir = os.path.join(output_dir,scene_name.split('.')[0])
        print('in:',scene_dir)
        # print('out:',save_dir)
        for j, video_name in tqdm(enumerate(os.listdir(scene_dir)),total=len(os.listdir(scene_dir))):
            
            # if j < 1000 :#or i >=2000:
                # print('index:',j)
                # continue
            print(video_name)
            video_path = os.path.join(scene_dir, video_name)
            if not video_path.endswith(".mp4"):# or os.path.isdir(output_dir):
                continue
            
            video_info = metadata[metadata['id'] == video_name[:-4]]
            num_frames = video_info['num frames'].iloc[0]

            scene_cam_dir = video_path.replace( "videos","annotations")[:-4]
            scene_cam_path = os.path.join(scene_cam_dir,'poses.npy')

            scene_caption_path = os.path.join(scene_cam_dir,'caption.json')

            with open(scene_caption_path, 'r', encoding='utf-8') as f:
                caption_data = json.load(f)
                caption = caption_data["SceneSummary"]
            if not os.path.exists(scene_cam_path):
                print(f"Pose not found: {scene_cam_path}")
                continue

            camera_poses = np.load(scene_cam_path)
            cam_data_len = camera_poses.shape[0]
                
                # cam_emb = {k: data[k].cpu() if isinstance(data[k], torch.Tensor) else data[k] for k in cam_data}
            # with open(scene_cam_path, 'rb') as f:
            #     cam_data = np.load(f)  # 此时cam_data仅包含数据，无文件句柄引用

            # 加载视频
            # video_path = scene_dir
            if not os.path.exists(video_path):
                print(f"Video not found: {video_path}")
                continue
            
            start_str = f"{0:07d}"
            end_str = f"{chunk_size:07d}"
            chunk_name = f"{video_name[:-4]}_{start_str}_{end_str}"
            first_save_chunk_dir = os.path.join(output_dir,chunk_name)
            
            first_chunk_encoded_path = os.path.join(first_save_chunk_dir, "encoded_video.pth")
            # print(first_chunk_encoded_path)
            if os.path.exists(first_chunk_encoded_path):
                data = torch.load(first_chunk_encoded_path,weights_only=False)
                if 'latents' in data:
                    video_frames = 1
            else:
                video_frames = encoder.load_video_frames(video_path)
                if video_frames is None:
                    print(f"Failed to load video: {video_path}")
                    continue
                print('video shape:',video_frames.shape)

                
                    
                video_frames = video_frames.unsqueeze(0).to("cuda", dtype=torch.bfloat16)
                print('video shape:',video_frames.shape)

            video_name = video_name[:-4].split('_')[0]
            start_frame = 0
            end_frame = num_frames
            # print("num_frames:",num_frames)

            cam_interval = end_frame // (cam_data_len - 1)
            
            cam_frames = np.linspace(start_frame, end_frame, cam_data_len, endpoint=True)
            cam_frames = np.round(cam_frames).astype(int)
            cam_frames = cam_frames.tolist()
            # list(range(0, end_frame + 1 , cam_interval))
            

            sampled_range = range(start_frame, end_frame , chunk_size)
            sampled_frames = list(sampled_range)

            sampled_chunk_end = sampled_frames[0] + chunk_size
            start_str = f"{sampled_frames[0]:07d}"
            end_str = f"{sampled_chunk_end:07d}"

            chunk_name = f"{video_name}_{start_str}_{end_str}"
            # save_chunk_path = os.path.join(output_dir,chunk_name,"encoded_video.pth")
            
            # if os.path.exists(save_chunk_path):
            #     print(f"Video {video_name} already encoded, skipping...")
            #     continue
            
            
            
            
            
            # print(sampled_frames)

            print(f"Encoding scene {video_name}...")
            chunk_count_in_one_video = 0
            for sampled_chunk_start in sampled_frames:
                if num_frames - sampled_chunk_start < 100:
                    continue
                sampled_chunk_end = sampled_chunk_start + chunk_size
                start_str = f"{sampled_chunk_start:07d}"
                end_str = f"{sampled_chunk_end:07d}"

                resample_cam_frame = list(range(sampled_chunk_start, sampled_chunk_end , 4))

                # 生成保存目录名（假设video_name已定义）
                chunk_name = f"{video_name}_{start_str}_{end_str}"
                save_chunk_dir = os.path.join(output_dir,chunk_name)

                os.makedirs(save_chunk_dir,exist_ok=True)  
                print(f"Encoding chunk {chunk_name}...")

                encoded_path = os.path.join(save_chunk_dir, "encoded_video.pth")

                missing_keys = required_keys
                if os.path.exists(encoded_path):
                    print('error:',encoded_path)
                    data = torch.load(encoded_path,weights_only=False)
                    missing_keys = [key for key in required_keys if key not in data]
                    # print(missing_keys)
                    # print(f"Chunk {chunk_name} already encoded, skipping...")
                    if missing_keys:
                        print(f"警告: 文件中缺少以下必要元素: {missing_keys}")
                    if len(missing_keys) == 0 :
                        continue
                else:
                    print(f"警告: 缺少pth文件: {encoded_path}")
                    if not isinstance(video_frames, torch.Tensor):
                        
                        video_frames = encoder.load_video_frames(video_path)
                        if video_frames is None:
                            print(f"Failed to load video: {video_path}")
                            continue
                            
                        video_frames = video_frames.unsqueeze(0).to("cuda", dtype=torch.bfloat16)

                    print('video shape:',video_frames.shape)
                if "latents" in missing_keys:
                    chunk_frames = video_frames[:,:, sampled_chunk_start - start_frame : sampled_chunk_end - start_frame,...]
                    
                    # print('extrinsic:',cam_emb['extrinsic'].shape)
                    
                    # chunk_cam_emb ={'extrinsic':cam_emb['extrinsic'][sampled_chunk_start - start_frame : sampled_chunk_end - start_frame],
                    #                 'intrinsic':cam_emb['intrinsic']}

                    # print('chunk shape:',chunk_frames.shape)

                    with torch.no_grad():
                        latents = encoder.pipe.encode_video(chunk_frames, **encoder.tiler_kwargs)[0]
                else:
                    latents = data['latents']
                if "cam_emb" in missing_keys:  
                    cam_emb = interpolate_camera_poses(cam_frames, camera_poses,resample_cam_frame)
                    chunk_cam_emb ={'extrinsic':cam_emb}
                    print(f"视频长度:{chunk_size},重采样相机长度:{cam_emb.shape[0]}")
                else:
                    chunk_cam_emb = data['cam_emb']

                if "prompt_emb" in missing_keys:
                    # 编码文本
                    if chunk_count_in_one_video == 0:
                        print(caption)
                        with torch.no_grad():
                            prompt_emb = encoder.pipe.encode_prompt(caption)
                else:
                    prompt_emb = data['prompt_emb']
                    
                    #     del encoder.pipe.prompter
                    # pdb.set_trace()
                    # 保存编码结果
                encoded_data = {
                        "latents": latents.cpu(),
                        "prompt_emb": {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in prompt_emb.items()},
                        "cam_emb": chunk_cam_emb
                    }
                    # pdb.set_trace()
                torch.save(encoded_data, encoded_path)
                print(f"Saved encoded data: {encoded_path}")
                processed_chunk_count += 1
                chunk_count_in_one_video += 1

        processed_count += 1

        print("Encoded scene numebr:",processed_count)
        print("Encoded chunk numebr:",processed_chunk_count)

        # os.makedirs(save_dir,exist_ok=True)  
        # # 检查是否已编码
        # encoded_path = os.path.join(save_dir, "encoded_video.pth")
        # if os.path.exists(encoded_path):
        #     print(f"Scene {scene_name} already encoded, skipping...")
        #     continue
        
        # 加载场景信息

        
        
        # try:
        # print(f"Encoding scene {scene_name}...")
        
        # 加载和编码视频
        
        # 编码视频
        # with torch.no_grad():
        #     latents = encoder.pipe.encode_video(video_frames, **encoder.tiler_kwargs)[0]
            
        #     # 编码文本
        #     if processed_count == 0:
        #         print('encode prompt!!!')
        #         prompt_emb = encoder.pipe.encode_prompt("A video of a scene shot using a pedestrian's front camera while walking")
        #         del encoder.pipe.prompter
        #     # pdb.set_trace()
        #     # 保存编码结果
        #     encoded_data = {
        #         "latents": latents.cpu(),
        #         #"prompt_emb": {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in prompt_emb.items()},
        #         "cam_emb": cam_emb
        #     }
        #     # pdb.set_trace()
        #     torch.save(encoded_data, encoded_path)
        #     print(f"Saved encoded data: {encoded_path}")
        #     processed_count += 1
                
        # except Exception as e:
        #     print(f"Error encoding scene {scene_name}: {e}")
        #     continue
    
    print(f"Encoding completed! Processed {processed_count} scenes.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--scenes_path", type=str, default="/share_zhuyixuan05/public_datasets/SpatialVID-HQ/SpatialVid/HQ/videos/")
    parser.add_argument("--text_encoder_path", type=str, 
                       default="models/Wan-AI/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth")
    parser.add_argument("--vae_path", type=str,
                       default="models/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth")
    
    parser.add_argument("--output_dir",type=str,
                       default="/share_zhuyixuan05/zhuyixuan05/spatialvid")

    args = parser.parse_args()
    encode_scenes(args.scenes_path, args.text_encoder_path, args.vae_path,args.output_dir)