scripts/encode_spatialvid_first_frame.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.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
        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_single_frame(self, video_path, frame_idx):
        """只加载指定的单帧"""
        reader = imageio.get_reader(video_path)
        
        try:
            # 直接跳转到指定帧
            frame_data = reader.get_data(frame_idx)
            frame = Image.fromarray(frame_data)
            frame = self.crop_and_resize(frame)
            frame = self.frame_process(frame)
            
            # 添加batch和time维度: [C, H, W] -> [1, C, 1, H, W]
            frame = frame.unsqueeze(0).unsqueeze(2)
            
        except Exception as e:
            print(f"Error loading frame {frame_idx} from {video_path}: {e}")
            return None
        finally:
            reader.close()
        
        return frame

    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

    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

    for i, scene_name in enumerate(os.listdir(scenes_path)):
        if i < 2:
            continue
        print('group:',i)
        scene_dir = os.path.join(scenes_path, scene_name)
        
        print('in:',scene_dir)
        for j, video_name in tqdm(enumerate(os.listdir(scene_dir)),total=len(os.listdir(scene_dir))):
            print(video_name)
            video_path = os.path.join(scene_dir, video_name)
            if not video_path.endswith(".mp4"):
                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]

            if not os.path.exists(video_path):
                print(f"Video not found: {video_path}")
                continue

            video_name = video_name[:-4].split('_')[0]
            start_frame = 0
            end_frame = 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()

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

            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}"

                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}...")

                first_latent_path = os.path.join(save_chunk_dir, "first_latent.pth")
                
                if os.path.exists(first_latent_path):
                    print(f"First latent for chunk {chunk_name} already exists, skipping...")
                    continue

                # 只加载需要的那一帧
                first_frame_idx = sampled_chunk_start
                print(f"first_frame:{first_frame_idx}")
                first_frame = encoder.load_single_frame(video_path, first_frame_idx)
                
                if first_frame is None:
                    print(f"Failed to load frame {first_frame_idx} from: {video_path}")
                    continue
                
                first_frame = first_frame.to("cuda", dtype=torch.bfloat16)
                
                # 重复4次
                repeated_first_frame = first_frame.repeat(1, 1, 4, 1, 1)
                print(f"Repeated first frame shape: {repeated_first_frame.shape}")

                with torch.no_grad():
                    first_latents = encoder.pipe.encode_video(repeated_first_frame, **encoder.tiler_kwargs)[0]

                first_latent_data = {
                    "latents": first_latents.cpu(),
                }
                torch.save(first_latent_data, first_latent_path)
                print(f"Saved first latent: {first_latent_path}")

                processed_chunk_count += 1
                chunk_count_in_one_video += 1

        processed_count += 1
        print("Encoded scene number:", processed_count)
        print("Encoded chunk number:", processed_chunk_count)
    
    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)