421 lines
17 KiB
Python
421 lines
17 KiB
Python
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, "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 = [self.parse_matrix(cam_data[f"frame{idx}"][f"cam{int(camera_type_map[camera_type]):02d}"]) for idx in cam_idx]
|
|
|
|
def generate_orbit_180(num_frames=81):
|
|
camera_data = {}
|
|
radius = 100 # Distance from center during orbit
|
|
center = np.array([3390, 1380, 240]) # Center point of orbit
|
|
|
|
# Arc movement similar to arc_left/arc_right but spanning 180 degrees
|
|
for i in range(num_frames):
|
|
# Calculate angle from 0 to 180 degrees
|
|
angle = i * 90.0 / (num_frames - 1)
|
|
angle_rad = np.radians(angle)
|
|
|
|
# Calculate position - circular path around center
|
|
x = center[0] + radius * np.cos(angle_rad)
|
|
y = center[1] + radius * np.sin(angle_rad)
|
|
z = center[2] # Z stays constant
|
|
pos = np.array([x, y, z])
|
|
|
|
# Calculate direction from camera to center point
|
|
dir_to_center = center - pos
|
|
dir_to_center_xy = dir_to_center.copy()
|
|
dir_to_center_xy[2] = 0 # Project to XY plane for horizontal orientation
|
|
|
|
# Normalize the direction vector
|
|
dir_to_center_xy = dir_to_center_xy / np.linalg.norm(dir_to_center_xy)
|
|
|
|
# For camera to face center, forward vector should be this direction
|
|
# The rotation matrix needs to make the camera's forward vector (local Z)
|
|
# point toward the center
|
|
|
|
# Calculate the angle between the camera and center in the XY plane
|
|
# This is the negative angle we need to rotate camera to face center
|
|
look_angle = np.arctan2(dir_to_center_xy[1], dir_to_center_xy[0])
|
|
|
|
# Rotation matrix for facing the center
|
|
cos_look = np.cos(look_angle)
|
|
sin_look = np.sin(look_angle)
|
|
|
|
# Format string with rotation that makes camera face the center
|
|
matrix_str = f"[{cos_look} {sin_look} 0 0] "
|
|
matrix_str += f"[{-sin_look} {cos_look} 0 0] "
|
|
matrix_str += f"[0 0 1 0] "
|
|
matrix_str += f"[{x} {y} {z} 1] "
|
|
|
|
# Store the camera string
|
|
frame_key = f"frame{i}"
|
|
if frame_key not in camera_data:
|
|
camera_data[frame_key] = {}
|
|
camera_data[frame_key]["cam11"] = matrix_str
|
|
return camera_data
|
|
|
|
# Generate the orbit camera data
|
|
#traj = [self.parse_matrix(generate_orbit_180(num_frames=81)[f"frame{idx}"][f"cam{int(11):02d}"]) for idx in cam_idx]
|
|
#print(traj)
|
|
|
|
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,)
|
|
|
|
class WanVideoSynCamMasterCameraEmbed:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Width of the image to encode"}),
|
|
"height": ("INT", {"default": 480, "min": 64, "max": 29048, "step": 8, "tooltip": "Height of the image to encode"}),
|
|
"num_frames": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 4, "tooltip": "Number of frames to encode"}),
|
|
|
|
"camera_type": ([
|
|
"azimuth",
|
|
"elevation",
|
|
"distance",
|
|
], {"tooltip": "Camera type to use"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("WANVIDIMAGE_EMBEDS", "CAMERAPOSES",)
|
|
RETURN_NAMES = ("camera_embeds", "camera_poses",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "WanVideoWrapper"
|
|
DESCRIPTION = "https://github.com/KwaiVGI/ReCamMaster"
|
|
|
|
def process(self, width, height, num_frames, camera_type):
|
|
# load camera
|
|
import json
|
|
from einops import rearrange
|
|
|
|
camera_data_path = os.path.join(script_directory, "syncam_extrinsics.json")
|
|
with open(camera_data_path, 'r') as file:
|
|
cam_data = json.load(file)
|
|
|
|
# load camera
|
|
multiview_c2ws = []
|
|
cam_idx = list(range(num_frames))[::4]
|
|
|
|
traj_list = []
|
|
|
|
if camera_type == "azimuth":
|
|
tgt_idx = 1
|
|
cond_idx = 3
|
|
elif camera_type == "elevation":
|
|
tgt_idx = 3
|
|
cond_idx = 7
|
|
elif camera_type == "distance":
|
|
tgt_idx = 9
|
|
cond_idx = 10
|
|
for view_idx in [cond_idx, tgt_idx]:
|
|
traj = [self.parse_matrix(cam_data[f"frame{idx}"][f"cam{view_idx:02d}"]) for idx in cam_idx]
|
|
traj = np.stack(traj).transpose(0, 2, 1)
|
|
traj_list.append(traj)
|
|
|
|
c2ws = []
|
|
for c2w in traj:
|
|
c2w = c2w[:, [1, 2, 0, 3]]
|
|
c2w[:3, 1] *= -1.
|
|
c2w[:3, 3] /= 200
|
|
c2ws.append(c2w)
|
|
multiview_c2ws.append(c2ws)
|
|
cond_cam_params = [Camera(cam_param) for cam_param in multiview_c2ws[0]]
|
|
tgt_cam_params = [Camera(cam_param) for cam_param in multiview_c2ws[1]]
|
|
relative_poses = []
|
|
for i in range(len(tgt_cam_params)):
|
|
relative_pose = self.get_relative_pose([tgt_cam_params[i], cond_cam_params[i]])
|
|
relative_poses.append(torch.as_tensor(relative_pose)[:,:3,:])
|
|
pose_embedding = torch.stack(relative_poses, dim=1) # v,21,3,4
|
|
pose_embedding = rearrange(pose_embedding, 'v f c d -> v f (c d)')
|
|
|
|
vae_stride = (4, 8, 8)
|
|
target_shape = (16, (num_frames - 1) // vae_stride[0] + 1,
|
|
height // vae_stride[1],
|
|
width // vae_stride[2])
|
|
|
|
embeds = {
|
|
"target_shape": target_shape,
|
|
"num_frames": num_frames,
|
|
"syncammaster": {
|
|
"camera_embed": pose_embedding,
|
|
}
|
|
}
|
|
|
|
return (embeds, np.concatenate(traj_list),)
|
|
|
|
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,
|
|
"WanVideoSynCamMasterCameraEmbed": WanVideoSynCamMasterCameraEmbed,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"WanVideoReCamMasterCameraEmbed": "WanVideo ReCamMaster Camera Embed",
|
|
"ReCamMasterPoseVisualizer": "ReCamMaster Pose Visualizer",
|
|
"WanVideoSynCamMasterCameraEmbed": "WanVideo SyncamMaster Camera Embed",
|
|
}
|