import os import sys import math from tqdm import tqdm from PIL import Image, ImageDraw project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) try: sys.path.append(os.path.join(project_root, "submodules/MoGe")) os.environ["TOKENIZERS_PARALLELISM"] = "false" except: print("Warning: MoGe not found, motion transfer will not be applied") import torch import numpy as np from PIL import Image import torchvision.transforms as transforms from diffusers import FluxControlPipeline, CogVideoXDPMScheduler from diffusers.utils import export_to_video, load_image, load_video from models.spatracker.predictor import SpaTrackerPredictor from models.spatracker.utils.visualizer import Visualizer from models.cogvideox_tracking import CogVideoXImageToVideoPipelineTracking from submodules.MoGe.moge.model import MoGeModel from image_gen_aux import DepthPreprocessor from moviepy.editor import ImageSequenceClip class DiffusionAsShaderPipeline: def __init__(self, gpu_id=0, output_dir='outputs'): """Initialize MotionTransfer class Args: gpu_id (int): GPU device ID output_dir (str): Output directory path """ # video parameters self.max_depth = 65.0 self.fps = 8 # camera parameters self.camera_motion=None self.fov=55 # device self.device = f"cuda:{gpu_id}" torch.cuda.set_device(gpu_id) self.dtype = torch.bfloat16 # files self.output_dir = output_dir os.makedirs(output_dir, exist_ok=True) # Initialize transform self.transform = transforms.Compose([ transforms.Resize((480, 720)), transforms.ToTensor() ]) @torch.no_grad() def _infer( self, prompt: str, model_path: str, tracking_tensor: torch.Tensor = None, image_tensor: torch.Tensor = None, # [C,H,W] in range [0,1] output_path: str = "./output.mp4", num_inference_steps: int = 25, guidance_scale: float = 6.0, num_videos_per_prompt: int = 1, dtype: torch.dtype = torch.bfloat16, fps: int = 24, seed: int = 42, ): """ Generates a video based on the given prompt and saves it to the specified path. Parameters: - prompt (str): The description of the video to be generated. - model_path (str): The path of the pre-trained model to be used. - tracking_tensor (torch.Tensor): Tracking video tensor [T, C, H, W] in range [0,1] - image_tensor (torch.Tensor): Input image tensor [C, H, W] in range [0,1] - output_path (str): The path where the generated video will be saved. - num_inference_steps (int): Number of steps for the inference process. - guidance_scale (float): The scale for classifier-free guidance. - num_videos_per_prompt (int): Number of videos to generate per prompt. - dtype (torch.dtype): The data type for computation. - seed (int): The seed for reproducibility. """ from transformers import T5EncoderModel, T5Tokenizer from diffusers import AutoencoderKLCogVideoX, CogVideoXDDIMScheduler from models.cogvideox_tracking import CogVideoXTransformer3DModelTracking vae = AutoencoderKLCogVideoX.from_pretrained(model_path, subfolder="vae") text_encoder = T5EncoderModel.from_pretrained(model_path, subfolder="text_encoder") tokenizer = T5Tokenizer.from_pretrained(model_path, subfolder="tokenizer") transformer = CogVideoXTransformer3DModelTracking.from_pretrained(model_path, subfolder="transformer") scheduler = CogVideoXDDIMScheduler.from_pretrained(model_path, subfolder="scheduler") pipe = CogVideoXImageToVideoPipelineTracking( vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=transformer, scheduler=scheduler ) # Convert tensor to PIL Image image_np = (image_tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8) image = Image.fromarray(image_np) height, width = image.height, image.width pipe.transformer.eval() pipe.text_encoder.eval() pipe.vae.eval() self.dtype = dtype # Process tracking tensor tracking_maps = tracking_tensor.float() # [T, C, H, W] tracking_maps = tracking_maps.to(device=self.device, dtype=dtype) tracking_first_frame = tracking_maps[0:1] # Get first frame as [1, C, H, W] height, width = tracking_first_frame.shape[2], tracking_first_frame.shape[3] # 2. Set Scheduler. pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") pipe.to(self.device, dtype=dtype) # pipe.enable_sequential_cpu_offload() pipe.vae.enable_slicing() pipe.vae.enable_tiling() pipe.transformer.eval() pipe.text_encoder.eval() pipe.vae.eval() pipe.transformer.gradient_checkpointing = False print("Encoding tracking maps") tracking_maps = tracking_maps.unsqueeze(0) # [B, T, C, H, W] tracking_maps = tracking_maps.permute(0, 2, 1, 3, 4) # [B, C, T, H, W] tracking_latent_dist = pipe.vae.encode(tracking_maps).latent_dist tracking_maps = tracking_latent_dist.sample() * pipe.vae.config.scaling_factor tracking_maps = tracking_maps.permute(0, 2, 1, 3, 4) # [B, F, C, H, W] # 4. Generate the video frames based on the prompt. video_generate = pipe( prompt=prompt, negative_prompt="The video is not of a high quality, it has a low resolution. Watermark present in each frame. The background is solid. Strange body and strange trajectory. Distortion.", image=image, num_videos_per_prompt=num_videos_per_prompt, num_inference_steps=num_inference_steps, num_frames=49, use_dynamic_cfg=True, guidance_scale=guidance_scale, generator=torch.Generator().manual_seed(seed), tracking_maps=tracking_maps, tracking_image=tracking_first_frame, height=height, width=width, ).frames[0] # 5. Export the generated frames to a video file. fps must be 8 for original video. output_path = output_path if output_path else f"result.mp4" os.makedirs(os.path.dirname(output_path), exist_ok=True) export_to_video(video_generate, output_path, fps=fps) #========== camera parameters ==========# def _set_camera_motion(self, camera_motion): self.camera_motion = camera_motion ##============= SpatialTracker =============## def generate_tracking_spatracker(self, video_tensor, density=70): """Generate tracking video Args: video_tensor (torch.Tensor): Input video tensor Returns: str: Path to tracking video """ print("Loading tracking models...") # Load tracking model tracker = SpaTrackerPredictor( checkpoint=os.path.join(project_root, 'checkpoints/spaT_final.pth'), interp_shape=(384, 576), seq_length=12 ).to(self.device) # Load depth model self.depth_preprocessor = DepthPreprocessor.from_pretrained("Intel/zoedepth-nyu-kitti") self.depth_preprocessor.to(self.device) try: video = video_tensor.unsqueeze(0).to(self.device) video_depths = [] for i in range(video_tensor.shape[0]): frame = (video_tensor[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) depth = self.depth_preprocessor(Image.fromarray(frame))[0] depth_tensor = transforms.ToTensor()(depth) # [1, H, W] video_depths.append(depth_tensor) video_depth = torch.stack(video_depths, dim=0).to(self.device) # print("Video depth shape:", video_depth.shape) segm_mask = np.ones((480, 720), dtype=np.uint8) pred_tracks, pred_visibility, T_Firsts = tracker( video * 255, video_depth=video_depth, grid_size=density, backward_tracking=False, depth_predictor=None, grid_query_frame=0, segm_mask=torch.from_numpy(segm_mask)[None, None].to(self.device), wind_length=12, progressive_tracking=False ) return pred_tracks.squeeze(0), pred_visibility.squeeze(0), T_Firsts finally: # Clean up GPU memory del tracker, self.depth_preprocessor torch.cuda.empty_cache() def visualize_tracking_spatracker(self, video, pred_tracks, pred_visibility, T_Firsts, save_tracking=True): video = video.unsqueeze(0).to(self.device) vis = Visualizer(save_dir=self.output_dir, grayscale=False, fps=24, pad_value=0) msk_query = (T_Firsts == 0) pred_tracks = pred_tracks[:,:,msk_query.squeeze()] pred_visibility = pred_visibility[:,:,msk_query.squeeze()] tracking_video = vis.visualize(video=video, tracks=pred_tracks, visibility=pred_visibility, save_video=False, filename="temp") tracking_video = tracking_video.squeeze(0) # [T, C, H, W] wide_list = list(tracking_video.unbind(0)) wide_list = [wide.permute(1, 2, 0).cpu().numpy() for wide in wide_list] clip = ImageSequenceClip(wide_list, fps=self.fps) tracking_path = None if save_tracking: try: tracking_path = os.path.join(self.output_dir, "tracking_video.mp4") clip.write_videofile(tracking_path, codec="libx264", fps=self.fps, logger=None) print(f"Video saved to {tracking_path}") except Exception as e: print(f"Warning: Failed to save tracking video: {e}") tracking_path = None # Convert tracking_video back to tensor in range [0,1] tracking_frames = np.array(list(clip.iter_frames())) / 255.0 tracking_video = torch.from_numpy(tracking_frames).permute(0, 3, 1, 2).float() return tracking_path, tracking_video ##============= MoGe =============## def valid_mask(self, pixels, W, H): """Check if pixels are within valid image bounds Args: pixels (numpy.ndarray): Pixel coordinates of shape [N, 2] W (int): Image width H (int): Image height Returns: numpy.ndarray: Boolean mask of valid pixels """ return ((pixels[:, 0] >= 0) & (pixels[:, 0] < W) & (pixels[:, 1] > 0) & \ (pixels[:, 1] < H)) def sort_points_by_depth(self, points, depths): """Sort points by depth values Args: points (numpy.ndarray): Points array of shape [N, 2] depths (numpy.ndarray): Depth values of shape [N] Returns: tuple: (sorted_points, sorted_depths, sort_index) """ # Combine points and depths into a single array for sorting combined = np.hstack((points, depths[:, None])) # Nx3 (points + depth) # Sort by depth (last column) in descending order sort_index = combined[:, -1].argsort()[::-1] sorted_combined = combined[sort_index] # Split back into points and depths sorted_points = sorted_combined[:, :-1] sorted_depths = sorted_combined[:, -1] return sorted_points, sorted_depths, sort_index def draw_rectangle(self, rgb, coord, side_length, color=(255, 0, 0)): """Draw a rectangle on the image Args: rgb (PIL.Image): Image to draw on coord (tuple): Center coordinates (x, y) side_length (int): Length of rectangle sides color (tuple): RGB color tuple """ draw = ImageDraw.Draw(rgb) # Calculate the bounding box of the rectangle left_up_point = (coord[0] - side_length//2, coord[1] - side_length//2) right_down_point = (coord[0] + side_length//2, coord[1] + side_length//2) color = tuple(list(color)) draw.rectangle( [left_up_point, right_down_point], fill=tuple(color), outline=tuple(color), ) def visualize_tracking_moge(self, points, mask, save_tracking=True): """Visualize tracking results from MoGe model Args: points (numpy.ndarray): Points array of shape [T, H, W, 3] mask (numpy.ndarray): Binary mask of shape [H, W] save_tracking (bool): Whether to save tracking video Returns: tuple: (tracking_path, tracking_video) - tracking_path (str): Path to saved tracking video, None if save_tracking is False - tracking_video (torch.Tensor): Tracking visualization tensor of shape [T, C, H, W] in range [0,1] """ # Create color array T, H, W, _ = points.shape colors = np.zeros((H, W, 3), dtype=np.uint8) # Set R channel - based on x coordinates (smaller on the left) colors[:, :, 0] = np.tile(np.linspace(0, 255, W), (H, 1)) # Set G channel - based on y coordinates (smaller on the top) colors[:, :, 1] = np.tile(np.linspace(0, 255, H), (W, 1)).T # Set B channel - based on depth z_values = points[0, :, :, 2] # get z values inv_z = 1 / z_values # calculate 1/z # Calculate 2% and 98% percentiles p2 = np.percentile(inv_z, 2) p98 = np.percentile(inv_z, 98) # Normalize to [0,1] range normalized_z = np.clip((inv_z - p2) / (p98 - p2), 0, 1) colors[:, :, 2] = (normalized_z * 255).astype(np.uint8) colors = colors.astype(np.uint8) points = points.reshape(T, -1, 3) colors = colors.reshape(-1, 3) # Initialize list to store frames frames = [] for i, pts_i in enumerate(tqdm(points, desc="rendering frames")): pixels, depths = pts_i[..., :2], pts_i[..., 2] pixels[..., 0] = pixels[..., 0] * W pixels[..., 1] = pixels[..., 1] * H pixels = pixels.astype(int) valid = self.valid_mask(pixels, W, H) frame_rgb = colors[valid] pixels = pixels[valid] depths = depths[valid] img = Image.fromarray(np.uint8(np.zeros([H, W, 3])), mode="RGB") sorted_pixels, _, sort_index = self.sort_points_by_depth(pixels, depths) step = 1 sorted_pixels = sorted_pixels[::step] sorted_rgb = frame_rgb[sort_index][::step] for j in range(sorted_pixels.shape[0]): self.draw_rectangle( img, coord=(sorted_pixels[j, 0], sorted_pixels[j, 1]), side_length=2, color=sorted_rgb[j], ) frames.append(np.array(img)) # Convert frames to video tensor in range [0,1] tracking_video = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).float() / 255.0 tracking_path = None if save_tracking: try: tracking_path = os.path.join(self.output_dir, "tracking_video_moge.mp4") # Convert back to uint8 for saving uint8_frames = [frame.astype(np.uint8) for frame in frames] clip = ImageSequenceClip(uint8_frames, fps=self.fps) clip.write_videofile(tracking_path, codec="libx264", fps=self.fps, logger=None) print(f"Video saved to {tracking_path}") except Exception as e: print(f"Warning: Failed to save tracking video: {e}") tracking_path = None return tracking_path, tracking_video ##============= CoTracker =============## def generate_tracking_cotracker(self, video_tensor, density=70): """Generate tracking video Args: video_tensor (torch.Tensor): Input video tensor Returns: tuple: (pred_tracks, pred_visibility) - pred_tracks (torch.Tensor): Tracking points with depth [T, N, 3] - pred_visibility (torch.Tensor): Visibility mask [T, N, 1] """ # Generate tracking points cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker3_offline").to(self.device) # Load depth model if not hasattr(self, 'depth_preprocessor') or self.depth_preprocessor is None: self.depth_preprocessor = DepthPreprocessor.from_pretrained("Intel/zoedepth-nyu-kitti") self.depth_preprocessor.to(self.device) try: video = video_tensor.unsqueeze(0).to(self.device) # Process all frames to get depth maps video_depths = [] for i in tqdm(range(video_tensor.shape[0]), desc="estimating depth"): frame = (video_tensor[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) depth = self.depth_preprocessor(Image.fromarray(frame))[0] depth_tensor = transforms.ToTensor()(depth) # [1, H, W] video_depths.append(depth_tensor) video_depth = torch.stack(video_depths, dim=0).to(self.device) # [T, 1, H, W] # Get tracking points and visibility print("tracking...") pred_tracks, pred_visibility = cotracker(video, grid_size=density) # B T N 2, B T N 1 # Extract dimensions B, T, N, _ = pred_tracks.shape H, W = video_depth.shape[2], video_depth.shape[3] # Create output tensor with depth pred_tracks_with_depth = torch.zeros((B, T, N, 3), device=self.device) pred_tracks_with_depth[:, :, :, :2] = pred_tracks # Copy x,y coordinates # Vectorized approach to get depths for all points # Reshape pred_tracks to process all batches and frames at once flat_tracks = pred_tracks.reshape(B*T, N, 2) # Clamp coordinates to valid image bounds x_coords = flat_tracks[:, :, 0].clamp(0, W-1).long() # [B*T, N] y_coords = flat_tracks[:, :, 1].clamp(0, H-1).long() # [B*T, N] # Get depths for all points at once # For each point in the flattened batch, get its depth from the corresponding frame depths = torch.zeros((B*T, N), device=self.device) for bt in range(B*T): t = bt % T # Time index depths[bt] = video_depth[t, 0, y_coords[bt], x_coords[bt]] # Reshape depths back to [B, T, N] and assign to output tensor pred_tracks_with_depth[:, :, :, 2] = depths.reshape(B, T, N) return pred_tracks_with_depth.squeeze(0), pred_visibility.squeeze(0) finally: del cotracker torch.cuda.empty_cache() def visualize_tracking_cotracker(self, points, vis_mask=None, save_tracking=True, point_wise=10, video_size=(480, 720)): """Visualize tracking results from CoTracker Args: points (torch.Tensor): Points array of shape [T, N, 3] vis_mask (torch.Tensor): Visibility mask of shape [T, N, 1] save_tracking (bool): Whether to save tracking video point_wise (int): Size of points in visualization video_size (tuple): Render size (height, width) Returns: tuple: (tracking_path, tracking_video) """ # Move tensors to CPU and convert to numpy if isinstance(points, torch.Tensor): points = points.detach().cpu().numpy() if vis_mask is not None and isinstance(vis_mask, torch.Tensor): vis_mask = vis_mask.detach().cpu().numpy() # Reshape if needed if vis_mask.ndim == 3 and vis_mask.shape[2] == 1: vis_mask = vis_mask.squeeze(-1) T, N, _ = points.shape H, W = video_size if vis_mask is None: vis_mask = np.ones((T, N), dtype=bool) colors = np.zeros((N, 3), dtype=np.uint8) first_frame_pts = points[0] u_min, u_max = 0, W u_normalized = np.clip((first_frame_pts[:, 0] - u_min) / (u_max - u_min), 0, 1) colors[:, 0] = (u_normalized * 255).astype(np.uint8) v_min, v_max = 0, H v_normalized = np.clip((first_frame_pts[:, 1] - v_min) / (v_max - v_min), 0, 1) colors[:, 1] = (v_normalized * 255).astype(np.uint8) z_values = first_frame_pts[:, 2] if np.all(z_values == 0): colors[:, 2] = np.random.randint(0, 256, N, dtype=np.uint8) else: inv_z = 1 / (z_values + 1e-10) p2 = np.percentile(inv_z, 2) p98 = np.percentile(inv_z, 98) normalized_z = np.clip((inv_z - p2) / (p98 - p2 + 1e-10), 0, 1) colors[:, 2] = (normalized_z * 255).astype(np.uint8) frames = [] for i in tqdm(range(T), desc="rendering frames"): pts_i = points[i] visibility = vis_mask[i] pixels, depths = pts_i[visibility, :2], pts_i[visibility, 2] pixels = pixels.astype(int) in_frame = self.valid_mask(pixels, W, H) pixels = pixels[in_frame] depths = depths[in_frame] frame_rgb = colors[visibility][in_frame] img = Image.fromarray(np.zeros((H, W, 3), dtype=np.uint8), mode="RGB") sorted_pixels, _, sort_index = self.sort_points_by_depth(pixels, depths) sorted_rgb = frame_rgb[sort_index] for j in range(sorted_pixels.shape[0]): self.draw_rectangle( img, coord=(sorted_pixels[j, 0], sorted_pixels[j, 1]), side_length=point_wise, color=sorted_rgb[j], ) frames.append(np.array(img)) # Convert frames to video tensor in range [0,1] tracking_video = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).float() / 255.0 tracking_path = None if save_tracking: try: tracking_path = os.path.join(self.output_dir, "tracking_video_cotracker.mp4") # Convert back to uint8 for saving uint8_frames = [frame.astype(np.uint8) for frame in frames] clip = ImageSequenceClip(uint8_frames, fps=self.fps) clip.write_videofile(tracking_path, codec="libx264", fps=self.fps, logger=None) print(f"Video saved to {tracking_path}") except Exception as e: print(f"Warning: Failed to save tracking video: {e}") tracking_path = None return tracking_path, tracking_video def apply_tracking(self, video_tensor, fps=8, tracking_tensor=None, img_cond_tensor=None, prompt=None, checkpoint_path=None, num_inference_steps=15): """Generate final video with motion transfer Args: video_tensor (torch.Tensor): Input video tensor [T,C,H,W] fps (float): Input video FPS tracking_tensor (torch.Tensor): Tracking video tensor [T,C,H,W] image_tensor (torch.Tensor): First frame tensor [C,H,W] to use for generation prompt (str): Generation prompt checkpoint_path (str): Path to model checkpoint """ self.fps = fps # Use first frame if no image provided if img_cond_tensor is None: img_cond_tensor = video_tensor[0] # Generate final video final_output = os.path.join(os.path.abspath(self.output_dir), "result.mp4") self._infer( prompt=prompt, model_path=checkpoint_path, tracking_tensor=tracking_tensor, image_tensor=img_cond_tensor, output_path=final_output, num_inference_steps=num_inference_steps, guidance_scale=6.0, dtype=torch.bfloat16, fps=self.fps ) print(f"Final video generated successfully at: {final_output}") def _set_object_motion(self, motion_type): """Set object motion type Args: motion_type (str): Motion direction ('up', 'down', 'left', 'right') """ self.object_motion = motion_type class FirstFrameRepainter: def __init__(self, gpu_id=0, output_dir='outputs'): """Initialize FirstFrameRepainter Args: gpu_id (int): GPU device ID output_dir (str): Output directory path """ self.device = f"cuda:{gpu_id}" self.output_dir = output_dir self.max_depth = 65.0 os.makedirs(output_dir, exist_ok=True) def repaint(self, image_tensor, prompt, depth_path=None, method="dav"): """Repaint first frame using Flux Args: image_tensor (torch.Tensor): Input image tensor [C,H,W] prompt (str): Repaint prompt depth_path (str): Path to depth image method (str): depth estimator, "moge" or "dav" or "zoedepth" Returns: torch.Tensor: Repainted image tensor [C,H,W] """ print("Loading Flux model...") # Load Flux model flux_pipe = FluxControlPipeline.from_pretrained( "black-forest-labs/FLUX.1-Depth-dev", torch_dtype=torch.bfloat16 ).to(self.device) # Get depth map if depth_path is None: if method == "moge": self.moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(self.device) depth_map = self.moge_model.infer(image_tensor.to(self.device))["depth"] depth_map = torch.clamp(depth_map, max=self.max_depth) depth_normalized = 1.0 - (depth_map / self.max_depth) depth_rgb = (depth_normalized * 255).cpu().numpy().astype(np.uint8) control_image = Image.fromarray(depth_rgb).convert("RGB") elif method == "zoedepth": self.depth_preprocessor = DepthPreprocessor.from_pretrained("Intel/zoedepth-nyu-kitti") self.depth_preprocessor.to(self.device) image_np = (image_tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8) control_image = self.depth_preprocessor(Image.fromarray(image_np))[0].convert("RGB") control_image = control_image.point(lambda x: 255 - x) # the zoedepth depth is inverted else: self.depth_preprocessor = DepthPreprocessor.from_pretrained("depth-anything/Depth-Anything-V2-Large-hf") self.depth_preprocessor.to(self.device) image_np = (image_tensor.permute(1, 2, 0).numpy() * 255).astype(np.uint8) control_image = self.depth_preprocessor(Image.fromarray(image_np))[0].convert("RGB") else: control_image = Image.open(depth_path).convert("RGB") try: repainted_image = flux_pipe( prompt=prompt, control_image=control_image, height=480, width=720, num_inference_steps=30, guidance_scale=7.5, ).images[0] # Save repainted image repainted_image.save(os.path.join(self.output_dir, "temp_repainted.png")) # Convert PIL Image to tensor transform = transforms.Compose([ transforms.ToTensor() ]) repainted_tensor = transform(repainted_image) return repainted_tensor finally: # Clean up GPU memory del flux_pipe if method == "moge": del self.moge_model else: del self.depth_preprocessor torch.cuda.empty_cache() class CameraMotionGenerator: def __init__(self, motion_type, frame_num=49, H=480, W=720, fx=None, fy=None, fov=55, device='cuda'): self.motion_type = motion_type self.frame_num = frame_num self.fov = fov self.device = device self.W = W self.H = H self.intr = torch.tensor([ [0, 0, W / 2], [0, 0, H / 2], [0, 0, 1] ], dtype=torch.float32, device=device) # if fx, fy not provided if not fx or not fy: fov_rad = math.radians(fov) fx = fy = (W / 2) / math.tan(fov_rad / 2) self.intr[0, 0] = fx self.intr[1, 1] = fy self.extr = torch.eye(4, device=device) def s2w_vggt(self, points, extrinsics, intrinsics): """ Transform points from pixel coordinates to world coordinates Args: points: Point cloud data of shape [T, N, 3] in uvz format extrinsics: Camera extrinsic matrices [B, T, 3, 4] or [T, 3, 4] intrinsics: Camera intrinsic matrices [B, T, 3, 3] or [T, 3, 3] Returns: world_points: Point cloud in world coordinates [T, N, 3] """ if isinstance(points, torch.Tensor): points = points.detach().cpu().numpy() if isinstance(extrinsics, torch.Tensor): extrinsics = extrinsics.detach().cpu().numpy() # Handle batch dimension if extrinsics.ndim == 4: # [B, T, 3, 4] extrinsics = extrinsics[0] # Take first batch if isinstance(intrinsics, torch.Tensor): intrinsics = intrinsics.detach().cpu().numpy() # Handle batch dimension if intrinsics.ndim == 4: # [B, T, 3, 3] intrinsics = intrinsics[0] # Take first batch T, N, _ = points.shape world_points = np.zeros_like(points) # Extract uvz coordinates uvz = points valid_mask = uvz[..., 2] > 0 # Create homogeneous coordinates [u, v, 1] uv_homogeneous = np.concatenate([uvz[..., :2], np.ones((T, N, 1))], axis=-1) # Transform from pixel to camera coordinates for i in range(T): K = intrinsics[i] K_inv = np.linalg.inv(K) R = extrinsics[i, :, :3] t = extrinsics[i, :, 3] R_inv = np.linalg.inv(R) valid_indices = np.where(valid_mask[i])[0] if len(valid_indices) > 0: valid_uv = uv_homogeneous[i, valid_indices] valid_z = uvz[i, valid_indices, 2] valid_xyz_camera = valid_uv @ K_inv.T valid_xyz_camera = valid_xyz_camera * valid_z[:, np.newaxis] # Transform from camera to world coordinates: X_world = R^-1 * (X_camera - t) valid_world_points = (valid_xyz_camera - t) @ R_inv.T world_points[i, valid_indices] = valid_world_points return world_points def w2s_vggt(self, world_points, extrinsics, intrinsics, poses=None): """ Project points from world coordinates to camera view Args: world_points: Point cloud in world coordinates [T, N, 3] extrinsics: Original camera extrinsic matrices [B, T, 3, 4] or [T, 3, 4] intrinsics: Camera intrinsic matrices [B, T, 3, 3] or [T, 3, 3] poses: Camera pose matrices [T, 4, 4], if None use first frame extrinsics Returns: camera_points: Point cloud in camera coordinates [T, N, 3] in uvz format """ if isinstance(world_points, torch.Tensor): world_points = world_points.detach().cpu().numpy() if isinstance(extrinsics, torch.Tensor): extrinsics = extrinsics.detach().cpu().numpy() if extrinsics.ndim == 4: extrinsics = extrinsics[0] if isinstance(intrinsics, torch.Tensor): intrinsics = intrinsics.detach().cpu().numpy() if intrinsics.ndim == 4: intrinsics = intrinsics[0] T, N, _ = world_points.shape # If no poses provided, use first frame extrinsics if poses is None: pose1 = np.eye(4) pose1[:3, :3] = extrinsics[0, :, :3] pose1[:3, 3] = extrinsics[0, :, 3] camera_poses = np.tile(pose1[np.newaxis, :, :], (T, 1, 1)) else: if isinstance(poses, torch.Tensor): camera_poses = poses.cpu().numpy() else: camera_poses = poses # Scale translation by 1/5 scaled_poses = camera_poses.copy() scaled_poses[:, :3, 3] = camera_poses[:, :3, 3] / 5.0 camera_poses = scaled_poses # Add homogeneous coordinates ones = np.ones([T, N, 1]) world_points_hom = np.concatenate([world_points, ones], axis=-1) # Transform points using batch matrix multiplication pts_cam_hom = np.matmul(world_points_hom, np.transpose(camera_poses, (0, 2, 1))) pts_cam = pts_cam_hom[..., :3] # Extract depth information depths = pts_cam[..., 2:3] valid_mask = depths[..., 0] > 0 # Normalize coordinates normalized_pts = pts_cam / (depths + 1e-10) # Apply intrinsic matrix for projection pts_pixel = np.matmul(normalized_pts, np.transpose(intrinsics, (0, 2, 1))) # Extract pixel coordinates u = pts_pixel[..., 0:1] v = pts_pixel[..., 1:2] # Set invalid points to zero u[~valid_mask] = 0 v[~valid_mask] = 0 depths[~valid_mask] = 0 # Return points in uvz format result = np.concatenate([u, v, depths], axis=-1) return torch.from_numpy(result) def w2s_moge(self, pts, poses): if isinstance(poses, np.ndarray): poses = torch.from_numpy(poses) assert poses.shape[0] == self.frame_num poses = poses.to(torch.float32).to(self.device) T, N, _ = pts.shape # (T, N, 3) intr = self.intr.unsqueeze(0).repeat(self.frame_num, 1, 1) ones = torch.ones((T, N, 1), device=self.device, dtype=pts.dtype) points_world_h = torch.cat([pts, ones], dim=-1) points_camera_h = torch.bmm(poses, points_world_h.permute(0, 2, 1)) points_camera = points_camera_h[:, :3, :].permute(0, 2, 1) points_image_h = torch.bmm(points_camera, intr.permute(0, 2, 1)) uv = points_image_h[:, :, :2] / points_image_h[:, :, 2:3] depth = points_camera[:, :, 2:3] # (T, N, 1) uvd = torch.cat([uv, depth], dim=-1) # (T, N, 3) return uvd def set_intr(self, K): if isinstance(K, np.ndarray): K = torch.from_numpy(K) self.intr = K.to(self.device) def set_extr(self, extr): if isinstance(extr, np.ndarray): extr = torch.from_numpy(extr) self.extr = extr.to(self.device) def rot_poses(self, angle, axis='y'): """Generate a single rotation matrix Args: angle (float): Rotation angle in degrees axis (str): Rotation axis ('x', 'y', or 'z') Returns: torch.Tensor: Single rotation matrix [4, 4] """ angle_rad = math.radians(angle) cos_theta = torch.cos(torch.tensor(angle_rad)) sin_theta = torch.sin(torch.tensor(angle_rad)) if axis == 'x': rot_mat = torch.tensor([ [1, 0, 0, 0], [0, cos_theta, -sin_theta, 0], [0, sin_theta, cos_theta, 0], [0, 0, 0, 1] ], dtype=torch.float32) elif axis == 'y': rot_mat = torch.tensor([ [cos_theta, 0, sin_theta, 0], [0, 1, 0, 0], [-sin_theta, 0, cos_theta, 0], [0, 0, 0, 1] ], dtype=torch.float32) elif axis == 'z': rot_mat = torch.tensor([ [cos_theta, -sin_theta, 0, 0], [sin_theta, cos_theta, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1] ], dtype=torch.float32) else: raise ValueError("Invalid axis value. Choose 'x', 'y', or 'z'.") return rot_mat.to(self.device) def trans_poses(self, dx, dy, dz): """ params: - dx: float, displacement along x axis。 - dy: float, displacement along y axis。 - dz: float, displacement along z axis。 ret: - matrices: torch.Tensor """ trans_mats = torch.eye(4).unsqueeze(0).repeat(self.frame_num, 1, 1) # (n, 4, 4) delta_x = dx / (self.frame_num - 1) delta_y = dy / (self.frame_num - 1) delta_z = dz / (self.frame_num - 1) for i in range(self.frame_num): trans_mats[i, 0, 3] = i * delta_x trans_mats[i, 1, 3] = i * delta_y trans_mats[i, 2, 3] = i * delta_z return trans_mats.to(self.device) def _look_at(self, camera_position, target_position): # look at direction direction = target_position - camera_position direction /= np.linalg.norm(direction) # calculate rotation matrix up = np.array([0, 1, 0]) right = np.cross(up, direction) right /= np.linalg.norm(right) up = np.cross(direction, right) rotation_matrix = np.vstack([right, up, direction]) rotation_matrix = np.linalg.inv(rotation_matrix) return rotation_matrix def spiral_poses(self, radius, forward_ratio = 0.5, backward_ratio = 0.5, rotation_times = 0.1, look_at_times = 0.5): """Generate spiral camera poses Args: radius (float): Base radius of the spiral forward_ratio (float): Scale factor for forward motion backward_ratio (float): Scale factor for backward motion rotation_times (float): Number of rotations to complete look_at_times (float): Scale factor for look-at point distance Returns: torch.Tensor: Camera poses of shape [num_frames, 4, 4] """ # Generate spiral trajectory t = np.linspace(0, 1, self.frame_num) r = np.sin(np.pi * t) * radius * rotation_times theta = 2 * np.pi * t # Calculate camera positions # Limit y motion for better floor/sky view y = r * np.cos(theta) * 0.3 x = r * np.sin(theta) z = -r z[z < 0] *= forward_ratio z[z > 0] *= backward_ratio # Set look-at target target_pos = np.array([0, 0, radius * look_at_times]) cam_pos = np.vstack([x, y, z]).T cam_poses = [] for pos in cam_pos: rot_mat = self._look_at(pos, target_pos) trans_mat = np.eye(4) trans_mat[:3, :3] = rot_mat trans_mat[:3, 3] = pos cam_poses.append(trans_mat[None]) camera_poses = np.concatenate(cam_poses, axis=0) return torch.from_numpy(camera_poses).to(self.device) def get_default_motion(self): """Parse motion parameters and generate corresponding motion matrices Supported formats: - trans [start_frame] [end_frame]: Translation motion - rot [start_frame] [end_frame]: Rotation motion - spiral [start_frame] [end_frame]: Spiral motion Multiple transformations can be combined using semicolon (;) as separator: e.g., "trans 0 0 0.5 0 30; rot x 25 0 30; trans 0.1 0 0 30 48" Note: - start_frame and end_frame are optional - frame range: 0-49 (will be clamped to this range) - if not specified, defaults to 0-49 - frames after end_frame will maintain the final transformation - for combined transformations, they are applied in sequence - moving left, up and zoom out is positive in video Returns: torch.Tensor: Motion matrices [num_frames, 4, 4] """ if not isinstance(self.motion_type, str): raise ValueError(f'camera_motion must be a string, but got {type(self.motion_type)}') # Split combined transformations transform_sequences = [s.strip() for s in self.motion_type.split(';')] # Initialize the final motion matrices final_motion = torch.eye(4, device=self.device).unsqueeze(0).repeat(49, 1, 1) # Process each transformation in sequence for transform in transform_sequences: params = transform.lower().split() if not params: continue motion_type = params[0] # Default frame range start_frame = 0 end_frame = 48 # 49 frames in total (0-48) if motion_type == 'trans': # Parse translation parameters if len(params) not in [4, 6]: raise ValueError(f"trans motion requires 3 or 5 parameters: 'trans ' or 'trans ', got: {transform}") dx, dy, dz = map(float, params[1:4]) if len(params) == 6: start_frame = max(0, min(48, int(params[4]))) end_frame = max(0, min(48, int(params[5]))) if start_frame > end_frame: start_frame, end_frame = end_frame, start_frame # Generate current transformation current_motion = torch.eye(4, device=self.device).unsqueeze(0).repeat(49, 1, 1) for frame_idx in range(49): if frame_idx < start_frame: continue elif frame_idx <= end_frame: t = (frame_idx - start_frame) / (end_frame - start_frame) current_motion[frame_idx, :3, 3] = torch.tensor([dx, dy, dz], device=self.device) * t else: current_motion[frame_idx] = current_motion[end_frame] # Combine with previous transformations final_motion = torch.matmul(final_motion, current_motion) elif motion_type == 'rot': # Parse rotation parameters if len(params) not in [3, 5]: raise ValueError(f"rot motion requires 2 or 4 parameters: 'rot ' or 'rot ', got: {transform}") axis = params[1] if axis not in ['x', 'y', 'z']: raise ValueError(f"Invalid rotation axis '{axis}', must be 'x', 'y' or 'z'") angle = float(params[2]) if len(params) == 5: start_frame = max(0, min(48, int(params[3]))) end_frame = max(0, min(48, int(params[4]))) if start_frame > end_frame: start_frame, end_frame = end_frame, start_frame current_motion = torch.eye(4, device=self.device).unsqueeze(0).repeat(49, 1, 1) for frame_idx in range(49): if frame_idx < start_frame: continue elif frame_idx <= end_frame: t = (frame_idx - start_frame) / (end_frame - start_frame) current_angle = angle * t current_motion[frame_idx] = self.rot_poses(current_angle, axis) else: current_motion[frame_idx] = current_motion[end_frame] # Combine with previous transformations final_motion = torch.matmul(final_motion, current_motion) elif motion_type == 'spiral': # Parse spiral motion parameters if len(params) not in [2, 4]: raise ValueError(f"spiral motion requires 1 or 3 parameters: 'spiral ' or 'spiral ', got: {transform}") radius = float(params[1]) if len(params) == 4: start_frame = max(0, min(48, int(params[2]))) end_frame = max(0, min(48, int(params[3]))) if start_frame > end_frame: start_frame, end_frame = end_frame, start_frame current_motion = torch.eye(4, device=self.device).unsqueeze(0).repeat(49, 1, 1) spiral_motion = self.spiral_poses(radius) for frame_idx in range(49): if frame_idx < start_frame: continue elif frame_idx <= end_frame: t = (frame_idx - start_frame) / (end_frame - start_frame) idx = int(t * (len(spiral_motion) - 1)) current_motion[frame_idx] = spiral_motion[idx] else: current_motion[frame_idx] = current_motion[end_frame] # Combine with previous transformations final_motion = torch.matmul(final_motion, current_motion) else: raise ValueError(f'camera_motion type must be in [trans, spiral, rot], but got {motion_type}') return final_motion class ObjectMotionGenerator: def __init__(self, device="cuda:0"): self.device = device self.num_frames = 49 def _get_points_in_mask(self, pred_tracks, mask): """Get points that lie within the mask Args: pred_tracks (torch.Tensor): Point trajectories [num_frames, num_points, 3] mask (torch.Tensor): Binary mask [H, W] Returns: torch.Tensor: Boolean mask for selected points [num_points] """ first_frame_points = pred_tracks[0] # [num_points, 3] xy_points = first_frame_points[:, :2] # [num_points, 2] xy_pixels = xy_points.round().long() xy_pixels[:, 0].clamp_(0, mask.shape[1] - 1) xy_pixels[:, 1].clamp_(0, mask.shape[0] - 1) points_in_mask = mask[xy_pixels[:, 1], xy_pixels[:, 0]] return points_in_mask def apply_motion(self, pred_tracks, mask, motion_type, distance, num_frames=49, tracking_method="spatracker"): self.num_frames = num_frames pred_tracks = pred_tracks.to(self.device).float() mask = mask.to(self.device) template = { 'up': ('trans', torch.tensor([0, -1, 0])), 'down': ('trans', torch.tensor([0, 1, 0])), 'left': ('trans', torch.tensor([-1, 0, 0])), 'right': ('trans', torch.tensor([1, 0, 0])), 'front': ('trans', torch.tensor([0, 0, 1])), 'back': ('trans', torch.tensor([0, 0, -1])), 'rot': ('rot', None) # rotate around y axis } if motion_type not in template: raise ValueError(f"unknown motion type: {motion_type}") motion_type, base_vec = template[motion_type] if base_vec is not None: base_vec = base_vec.to(self.device) * distance if tracking_method == "moge": T, H, W, _ = pred_tracks.shape valid_selected = ~torch.any(torch.isnan(pred_tracks[0]), dim=2) & mask points = pred_tracks[0][valid_selected].reshape(-1, 3) else: points_in_mask = self._get_points_in_mask(pred_tracks, mask) points = pred_tracks[0, points_in_mask] center = points.mean(dim=0) motions = [] for frame_idx in range(num_frames): t = frame_idx / (num_frames - 1) current_motion = torch.eye(4, device=self.device) current_motion[:3, 3] = -center motion_mat = torch.eye(4, device=self.device) if motion_type == 'trans': motion_mat[:3, 3] = base_vec * t else: # 'rot' angle_rad = torch.deg2rad(torch.tensor(distance * t, device=self.device)) cos_t = torch.cos(angle_rad) sin_t = torch.sin(angle_rad) motion_mat[0, 0] = cos_t motion_mat[0, 2] = sin_t motion_mat[2, 0] = -sin_t motion_mat[2, 2] = cos_t current_motion = motion_mat @ current_motion current_motion[:3, 3] += center motions.append(current_motion) motions = torch.stack(motions) # [num_frames, 4, 4] if tracking_method == "moge": modified_tracks = pred_tracks.clone().reshape(T, -1, 3) valid_selected = valid_selected.reshape([-1]) for frame_idx in range(self.num_frames): motion_mat = motions[frame_idx] if W > 1: motion_mat = motion_mat.clone() motion_mat[0, 3] /= W motion_mat[1, 3] /= H points = modified_tracks[frame_idx, valid_selected] points_homo = torch.cat([points, torch.ones_like(points[:, :1])], dim=1) transformed_points = torch.matmul(points_homo, motion_mat.T) modified_tracks[frame_idx, valid_selected] = transformed_points[:, :3] return modified_tracks.reshape(T, H, W, 3) else: points_in_mask = self._get_points_in_mask(pred_tracks, mask) modified_tracks = pred_tracks.clone() for frame_idx in range(pred_tracks.shape[0]): motion_mat = motions[frame_idx] points = modified_tracks[frame_idx, points_in_mask] points_homo = torch.cat([points, torch.ones_like(points[:, :1])], dim=1) transformed_points = torch.matmul(points_homo, motion_mat.T) modified_tracks[frame_idx, points_in_mask] = transformed_points[:, :3] return modified_tracks