diff --git a/__init__.py b/__init__.py index 2e96bd6..42c719f 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,7 @@ from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .recammaster.nodes import NODE_CLASS_MAPPINGS as RECAM_MASTER_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as RECAM_MASTER_NODE_DISPLAY_NAME_MAPPINGS + +NODE_CLASS_MAPPINGS.update(RECAM_MASTER_NODE_CLASS_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(RECAM_MASTER_NODE_DISPLAY_NAME_MAPPINGS) __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/nodes.py b/nodes.py index 49ff0e3..d86b72e 100644 --- a/nodes.py +++ b/nodes.py @@ -2033,122 +2033,6 @@ class WanVideoVACEStartToEndFrame: return (out_batch.cpu().float(), masks.cpu().float()) -#region ReCamMaster - -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) - - -class WanVideoReCamMasterCameraEmbed: - @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 = ("WANVIDIMAGE_EMBEDS", ) - RETURN_NAMES = ("camera_embeds",) - FUNCTION = "process" - CATEGORY = "WanVideoWrapper" - DESCRIPTION = "https://github.com/KwaiVGI/ReCamMaster" - - def process(self, camera_type, latents): - # load camera - import json - from einops import rearrange - - camera_data_path = os.path.join(script_directory, "camera_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] - print("cam_idx", cam_idx) - traj = [self.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) - c2ws = [] - for c2w in traj: - 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,) - - def parse_matrix(self, 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) - - 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 #region context options class WanVideoContextOptions: @@ -3357,7 +3241,6 @@ NODE_CLASS_MAPPINGS = { "WanVideoVACEEncode": WanVideoVACEEncode, "WanVideoVACEStartToEndFrame": WanVideoVACEStartToEndFrame, "WanVideoVACEModelSelect": WanVideoVACEModelSelect, - "WanVideoReCamMasterCameraEmbed": WanVideoReCamMasterCameraEmbed, } NODE_DISPLAY_NAME_MAPPINGS = { "WanVideoSampler": "WanVideo Sampler", @@ -3393,5 +3276,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "WanVideoVACEEncode": "WanVideo VACE Encode", "WanVideoVACEStartToEndFrame": "WanVideo VACE Start To End Frame", "WanVideoVACEModelSelect": "WanVideo VACE Model Select", - "WanVideoReCamMasterCameraEmbed": "WanVideo ReCam Master Camera Embed", } diff --git a/camera_extrinsics.json b/recammaster/camera_extrinsics.json similarity index 100% rename from camera_extrinsics.json rename to recammaster/camera_extrinsics.json diff --git a/recammaster/nodes.py b/recammaster/nodes.py new file mode 100644 index 0000000..de4f203 --- /dev/null +++ b/recammaster/nodes.py @@ -0,0 +1,281 @@ +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) + + +class WanVideoReCamMasterCameraEmbed: + @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 = ("WANVIDIMAGE_EMBEDS", "CAMERAPOSES",) + RETURN_NAMES = ("camera_embeds", "camera_poses",) + FUNCTION = "process" + CATEGORY = "WanVideoWrapper" + DESCRIPTION = "https://github.com/KwaiVGI/ReCamMaster" + + def process(self, camera_type, latents): + # load camera + import json + from einops import rearrange + + camera_data_path = os.path.join(script_directory, "camera_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 = [self.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) + c2ws = [] + for c2w in traj: + 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, traj,) + + def parse_matrix(self, 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) + + 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, + } +NODE_DISPLAY_NAME_MAPPINGS = { + "WanVideoReCamMasterCameraEmbed": "WanVideo ReCamMaster Camera Embed", + "ReCamMasterPoseVisualizer": "ReCamMaster Pose Visualizer", + }