recammaster/nodes.py

import numpy as np
import os, io
import torch
from PIL import Image
import math, time
script_directory = os.path.dirname(os.path.abspath(__file__))

class Camera(object):
    def __init__(self, c2w):
        c2w_mat = np.array(c2w).reshape(4, 4)
        self.c2w_mat = c2w_mat
        self.w2c_mat = np.linalg.inv(c2w_mat)

def parse_matrix(matrix_str):
    rows = matrix_str.strip().split('] [')
    matrix = []
    for row in rows:
        row = row.replace('[', '').replace(']', '')
        matrix.append(list(map(float, row.split())))
    return np.array(matrix)

class WanVideoReCamMasterDefaultCamera:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "camera_type": ([
                "pan_right", 
                "pan_left",
                "tilt_up",
                "tilt_down",
                "zoom_in",
                "zoom_out",
                "translate_up",
                "translate_down",
                "arc_left",
                "arc_right",
                ], {"default": "pan_right", "tooltip": "Camera type to use"}),
            "latents": ("LATENT", {"tooltip": "source video"}),
        },
        }

    RETURN_TYPES = ("CAMERAPOSES",)
    RETURN_NAMES = ("camera_poses",)
    FUNCTION = "process"
    CATEGORY = "WanVideoWrapper"
    DESCRIPTION = "https://github.com/KwaiVGI/ReCamMaster"

    def process(self, camera_type, latents):
        import json
        
        camera_data_path = os.path.join(script_directory, "recam_extrinsics.json")
        with open(camera_data_path, 'r') as file:
            cam_data = json.load(file)
        
        samples = latents["samples"].squeeze(0)
        C, T, H, W = samples.shape
        num_frames = (T - 1) * 4 + 1

        camera_type_map = {
            "pan_right": 1,
            "pan_left": 2,
            "tilt_up": 3,
            "tilt_down": 4,
            "zoom_in": 5,
            "zoom_out": 6,
            "translate_up": 7,
            "translate_down": 8,
            "arc_left": 9,
            "arc_right": 10,
        }

        cam_idx = list(range(num_frames))[::4]
        traj = [parse_matrix(cam_data[f"frame{idx}"][f"cam{int(camera_type_map[camera_type]):02d}"]) for idx in cam_idx]
        traj = np.stack(traj).transpose(0, 2, 1)
       
        return (traj,)
        
class WanVideoReCamMasterGenerateOrbitCamera:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "num_frames": ("INT", {"default": 81, "min": 1, "max": 1000, "step": 1, "tooltip": "Number of frames to generate"}),
            "degrees": ("INT", {"default": 90, "min": -180, "max": 180, "step": 1, "tooltip": "Degrees to orbit"}),
        },
        }

    RETURN_TYPES = ("CAMERAPOSES",)
    RETURN_NAMES = ("camera_poses",)
    FUNCTION = "process"
    CATEGORY = "WanVideoWrapper"
    DESCRIPTION = "https://github.com/KwaiVGI/ReCamMaster"

    def process(self, degrees, num_frames):

        def generate_orbit(num_frames=num_frames, degrees=degrees):
            camera_data = []
            center = np.array([3390, 1380, 240])  # Center point of orbit
            
            for i in range(num_frames):
                # Calculate angle from 0 to specified degrees
                angle = i * degrees / (num_frames - 1)
                angle_rad = np.radians(angle)
                
                # Calculate position - circular path around center
                x = center[0] - np.cos(angle_rad)
                y = center[1] - np.sin(angle_rad)
                z = center[2]
                
                # Calculate direction from camera to center point
                camera_pos = np.array([x, y, z])
                dir_to_center = center - camera_pos
                
                # Calculate the angle needed to face the center
                look_angle = np.arctan2(dir_to_center[1], dir_to_center[0])
                
                # Rotation matrix for facing the center (corrected)
                cos_look = np.cos(look_angle)
                sin_look = np.sin(look_angle)
                
                # Create transformation matrix directly
                transform = np.array([
                    [cos_look, -sin_look, 0, x],
                    [sin_look, cos_look, 0, y],
                    [0, 0, 1, z],
                    [0, 0, 0, 1]
                ])
                
                camera_data.append(transform)
                
            return camera_data
        
        # Generate orbit data
        camera_transforms = generate_orbit(num_frames=num_frames, degrees=degrees)
        
        traj = camera_transforms[::4]
        traj = np.stack(traj)
       
        return (traj,)
class WanVideoReCamMasterCameraEmbed:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "camera_poses": ("CAMERAPOSES",),
            "latents": ("LATENT", {"tooltip": "source video"}),
        },
        }

    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS", "CAMERAPOSES",)
    RETURN_NAMES = ("camera_embeds", "camera_poses",)
    FUNCTION = "process"
    CATEGORY = "WanVideoWrapper"
    DESCRIPTION = "https://github.com/KwaiVGI/ReCamMaster"

    def process(self, camera_poses, latents):
        from einops import rearrange
        samples = latents["samples"].squeeze(0)
        C, T, H, W = samples.shape
        num_frames = (T - 1) * 4 + 1
        
        
        c2ws = []
        for c2w in camera_poses:
            c2w = c2w[:, [1, 2, 0, 3]]
            c2w[:3, 1] *= -1.
            c2w[:3, 3] /= 100
            c2ws.append(c2w)
        tgt_cam_params = [Camera(cam_param) for cam_param in c2ws]
        relative_poses = []
        for i in range(len(tgt_cam_params)):
            relative_pose = self.get_relative_pose([tgt_cam_params[0], tgt_cam_params[i]])
            relative_poses.append(torch.as_tensor(relative_pose)[:,:3,:][1])
        pose_embedding = torch.stack(relative_poses, dim=0)  # 21x3x4
        pose_embedding = rearrange(pose_embedding, 'b c d -> b (c d)')

        seq_len = math.ceil((H * W) / 4 * ((num_frames - 1) // 4 + 1))
      
        embeds = {
            "max_seq_len": seq_len,
            "target_shape": samples.shape,
            "num_frames": num_frames,
            "recammaster": {
                "camera_embed": pose_embedding,
                "source_latents": samples
            }
        }

        return (embeds, camera_poses,)
    
    def get_relative_pose(self, cam_params):
        abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]
        abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]

        cam_to_origin = 0
        target_cam_c2w = np.array([
            [1, 0, 0, 0],
            [0, 1, 0, -cam_to_origin],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ])
        abs2rel = target_cam_c2w @ abs_w2cs[0]
        ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]
        ret_poses = np.array(ret_poses, dtype=np.float32)
        return ret_poses

def get_c2w(w2cs, transform_matrix, relative_c2w=True):
    if relative_c2w:
        target_cam_c2w = np.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ])
        abs2rel = target_cam_c2w @ w2cs[0]
        ret_poses = [target_cam_c2w, ] + [abs2rel @ np.linalg.inv(w2c) for w2c in w2cs[1:]]
    else:
        ret_poses = [np.linalg.inv(w2c) for w2c in w2cs]
    ret_poses = [transform_matrix @ x for x in ret_poses]
    return np.array(ret_poses, dtype=np.float32)

class ReCamMasterPoseVisualizer:
                
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "camera_poses": ("CAMERAPOSES",),
            "base_xval": ("FLOAT", {"default": 0.2,"min": 0, "max": 100, "step": 0.01}),
            "zval": ("FLOAT", {"default": 0.3,"min": 0, "max": 100, "step": 0.01}),
            "scale": ("FLOAT", {"default": 1.0,"min": 0.01, "max": 10.0, "step": 0.01}),
            "arrow_length": ("FLOAT", {"default": 1,"min": 0, "max": 100, "step": 0.01}),
            },
            }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "plot"
    CATEGORY = "WanVideoWrapper"
    DESCRIPTION = """
Visualizes the camera poses, from Animatediff-Evolved CameraCtrl Pose  
or a .txt file with RealEstate camera intrinsics and coordinates, in a 3D plot. 
"""
        
    def plot(self, camera_poses, scale, base_xval, zval, arrow_length):
        import matplotlib as mpl
        mpl.use('Agg')
        
        import matplotlib.pyplot as plt
        from torchvision.transforms import ToTensor

        x_min = -2.0 * scale
        x_max = 2.0 * scale
        y_min = -2.0 * scale
        y_max = 2.0 * scale
        z_min = -2.0 * scale
        z_max = 2.0 * scale
        plt.rcParams['text.color'] = '#999999'
        self.fig = plt.figure(figsize=(18, 7))
        self.fig.patch.set_facecolor('#353535')
        self.ax = self.fig.add_subplot(projection='3d')
        self.ax.set_facecolor('#353535') # Set the background color here
        self.ax.grid(color='#999999', linestyle='-', linewidth=0.5)
        self.plotly_data = None  # plotly data traces
        self.ax.set_aspect("auto")
        self.ax.set_xlim(x_min, x_max)
        self.ax.set_ylim(y_min, y_max)
        self.ax.set_zlim(z_min, z_max)
        self.ax.set_xlabel('x', color='#999999')
        self.ax.set_ylabel('y', color='#999999')
        self.ax.set_zlabel('z', color='#999999')
        for text in self.ax.get_xticklabels() + self.ax.get_yticklabels() + self.ax.get_zticklabels():
            text.set_color('#999999')
        print('initialize camera pose visualizer')

        total_frames = len(camera_poses)

        w2cs = []
        for cam in camera_poses:
            if cam.shape[0] == 3:
                cam = np.vstack((cam, np.array([[0, 0, 0, 1]])))
            cam = cam[:, [1, 2, 0, 3]]
            cam[:3, 1] *= -1.
            w2cs.append(np.linalg.inv(cam))
        transform_matrix = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]])
        c2ws = get_c2w(w2cs, transform_matrix, True)
        scale = max(max(abs(c2w[:3, 3])) for c2w in c2ws)
        if scale > 1e-3:  # otherwise, pan or tilt
            for c2w in c2ws:
                c2w[:3, 3] /= scale

        for frame_idx, c2w in enumerate(c2ws):
            self.extrinsic2pyramid(c2w, frame_idx / total_frames, hw_ratio=1, base_xval=base_xval, zval=(zval))
          
            if arrow_length > 0:
                pos = c2w[:3, 3]
                forward = c2w[:3, 2]
                arrow_start = pos + forward * base_xval
                arrow_length = arrow_length
                self.ax.quiver(arrow_start[0], arrow_start[1], arrow_start[2], 
                    forward[0], forward[1], forward[2], 
                    color='black', length=arrow_length, arrow_length_ratio=0.1)

        # Create the colorbar
        cmap = mpl.cm.rainbow
        norm = mpl.colors.Normalize(vmin=0, vmax=total_frames)
        colorbar = self.fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), ax=self.ax, orientation='vertical')

        # Change the colorbar label
        colorbar.set_label('Frame', color='#999999') # Change the label and its color

        # Change the tick colors
        colorbar.ax.yaxis.set_tick_params(colors='#999999') # Change the tick color

        # Change the tick frequency
        # Assuming you want to set the ticks at every 10th frame
        ticks = np.arange(0, total_frames, 10)
        colorbar.ax.yaxis.set_ticks(ticks)
        
        plt.title('')
        plt.draw()
        buf = io.BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
        buf.seek(0)
        img = Image.open(buf)
        tensor_img = ToTensor()(img)
        buf.close()
        tensor_img = tensor_img.permute(1, 2, 0).unsqueeze(0)
        return (tensor_img,)

    def extrinsic2pyramid(self, extrinsic, color_map='red', hw_ratio=9/16, base_xval=1, zval=3):
        from mpl_toolkits.mplot3d.art3d import Poly3DCollection
        import matplotlib.pyplot as plt
        vertex_std = np.array([[0, 0, 0, 1],
                               [base_xval, -base_xval * hw_ratio, zval, 1],
                               [base_xval, base_xval * hw_ratio, zval, 1],
                               [-base_xval, base_xval * hw_ratio, zval, 1],
                               [-base_xval, -base_xval * hw_ratio, zval, 1]])
        vertex_transformed = vertex_std @ extrinsic.T
        meshes = [[vertex_transformed[0, :-1], vertex_transformed[1][:-1], vertex_transformed[2, :-1]],
                            [vertex_transformed[0, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1]],
                            [vertex_transformed[0, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]],
                            [vertex_transformed[0, :-1], vertex_transformed[4, :-1], vertex_transformed[1, :-1]],
                            [vertex_transformed[1, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]]]

        color = color_map if isinstance(color_map, str) else plt.cm.rainbow(color_map)

        self.ax.add_collection3d(
            Poly3DCollection(meshes, facecolors=color, linewidths=0.3, edgecolors=color, alpha=0.35))

    def customize_legend(self, list_label):
        from matplotlib.patches import Patch
        import matplotlib.pyplot as plt
        list_handle = []
        for idx, label in enumerate(list_label):
            color = plt.cm.rainbow(idx / len(list_label))
            patch = Patch(color=color, label=label)
            list_handle.append(patch)
        plt.legend(loc='right', bbox_to_anchor=(1.8, 0.5), handles=list_handle)

NODE_CLASS_MAPPINGS = {
    "WanVideoReCamMasterCameraEmbed": WanVideoReCamMasterCameraEmbed,
    "ReCamMasterPoseVisualizer": ReCamMasterPoseVisualizer,
    "WanVideoReCamMasterGenerateOrbitCamera": WanVideoReCamMasterGenerateOrbitCamera,
    "WanVideoReCamMasterDefaultCamera": WanVideoReCamMasterDefaultCamera,
    }
NODE_DISPLAY_NAME_MAPPINGS = {
    "WanVideoReCamMasterCameraEmbed": "WanVideo ReCamMaster Camera Embed",
    "ReCamMasterPoseVisualizer": "ReCamMaster Pose Visualizer",
    "WanVideoReCamMasterGenerateOrbitCamera": "WanVideo ReCamMaster Generate Orbit Camera",
    "WanVideoReCamMasterDefaultCamera": "WanVideo ReCamMaster Default Camera",
    }