from typing import Callable, Dict, List, Optional, Union import numpy as np import PIL.Image import torch import torch.nn.functional as F import torchvision.transforms.functional as TF import kornia from matplotlib import cm from torchvision.io import write_video from PIL import Image, ImageOps import os from typing import Union, Tuple, List import math from matplotlib import pyplot as plt from mpl_toolkits.mplot3d.art3d import Poly3DCollection DEFAULT_FOV_RAD = 0.9424777960769379 # 54 degrees by default def get_default_intrinsics( fov_rad=DEFAULT_FOV_RAD, aspect_ratio=1.0, ): if not isinstance(fov_rad, torch.Tensor): fov_rad = torch.tensor( [fov_rad] if isinstance(fov_rad, (int, float)) else fov_rad ) if aspect_ratio >= 1.0: # W >= H focal_x = 0.5 / torch.tan(0.5 * fov_rad) focal_y = focal_x * aspect_ratio else: # W < H focal_y = 0.5 / torch.tan(0.5 * fov_rad) focal_x = focal_y / aspect_ratio intrinsics = focal_x.new_zeros((focal_x.shape[0], 3, 3)) intrinsics[:, torch.eye(3, device=focal_x.device, dtype=bool)] = torch.stack( [focal_x, focal_y, torch.ones_like(focal_x)], dim=-1 ) intrinsics[:, :, -1] = torch.tensor( [0.5, 0.5, 1.0], device=focal_x.device, dtype=focal_x.dtype ) return intrinsics def to_hom(X): # get homogeneous coordinates of the input X_hom = torch.cat([X, torch.ones_like(X[..., :1])], dim=-1) return X_hom def to_hom_pose(pose): # get homogeneous coordinates of the input pose if pose.shape[-2:] == (3, 4): pose_hom = torch.eye(4, device=pose.device)[None].repeat(pose.shape[0], 1, 1) pose_hom[:, :3, :] = pose return pose_hom return pose def get_image_grid(img_h, img_w): # add 0.5 is VERY important especially when your img_h and img_w # is not very large (e.g., 72)!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! y_range = torch.arange(img_h, dtype=torch.float32).add_(0.5) x_range = torch.arange(img_w, dtype=torch.float32).add_(0.5) Y, X = torch.meshgrid(y_range, x_range, indexing="ij") # [H,W] xy_grid = torch.stack([X, Y], dim=-1).view(-1, 2) # [HW,2] return to_hom(xy_grid) # [HW,3] def img2cam(X, cam_intr): return X @ cam_intr.inverse().transpose(-1, -2) def cam2world(X, pose): X_hom = to_hom(X) pose_inv = torch.linalg.inv(to_hom_pose(pose))[..., :3, :4] return X_hom @ pose_inv.transpose(-1, -2) def get_center_and_ray(img_h, img_w, pose, intr): # [HW,2] # given the intrinsic/extrinsic matrices, get the camera center and ray directions] # assert(opt.camera.model=="perspective") # compute center and ray grid_img = get_image_grid(img_h, img_w) # [HW,3] grid_3D_cam = img2cam(grid_img.to(intr.device), intr.float()) # [B,HW,3] center_3D_cam = torch.zeros_like(grid_3D_cam) # [B,HW,3] # transform from camera to world coordinates grid_3D = cam2world(grid_3D_cam, pose) # [B,HW,3] center_3D = cam2world(center_3D_cam, pose) # [B,HW,3] ray = grid_3D - center_3D # [B,HW,3] return center_3D, ray, grid_3D_cam def get_plucker_coordinates( extrinsics_src, extrinsics, intrinsics=None, fov_rad=DEFAULT_FOV_RAD, target_size=[72, 72], ): # Support for batch dimension has_batch_dim = len(extrinsics.shape) == 4 if has_batch_dim: # [B, N, 4, 4] -> reshape to handle batch batch_size, num_cameras = extrinsics.shape[0:2] extrinsics_flat = extrinsics.reshape(-1, *extrinsics.shape[2:]) # Handle extrinsics_src appropriately if len(extrinsics_src.shape) == 3: # [B, 4, 4] extrinsics_src_expanded = extrinsics_src.unsqueeze(1).expand(-1, num_cameras, -1, -1) extrinsics_src_flat = extrinsics_src_expanded.reshape(-1, *extrinsics_src.shape[1:]) else: # [4, 4] - single extrinsics_src for all batches extrinsics_src_flat = extrinsics_src.expand(batch_size * num_cameras, -1, -1) # Handle intrinsics for batch if intrinsics is None: intrinsics = get_default_intrinsics(fov_rad).to(extrinsics.device) intrinsics = intrinsics.expand(batch_size * num_cameras, -1, -1) elif len(intrinsics.shape) == 3: # [N, 3, 3] if intrinsics.shape[0] == num_cameras: intrinsics = intrinsics.expand(batch_size, -1, -1, -1).reshape(-1, *intrinsics.shape[1:]) else: intrinsics = intrinsics.expand(batch_size * num_cameras, -1, -1) elif len(intrinsics.shape) == 4: # [B, N, 3, 3] intrinsics = intrinsics.reshape(-1, *intrinsics.shape[2:]) else: # Original behavior for non-batch input extrinsics_flat = extrinsics extrinsics_src_flat = extrinsics_src if intrinsics is None: intrinsics = get_default_intrinsics(fov_rad).to(extrinsics.device) # Process intrinsics normalization if not ( torch.all(intrinsics[:, :2, -1] >= 0) and torch.all(intrinsics[:, :2, -1] <= 1) ): intrinsics[:, :2] /= intrinsics.new_tensor(target_size).view(1, -1, 1) * 8 # Ensure normalized intrinsics assert ( torch.all(intrinsics[:, :2, -1] >= 0) and torch.all(intrinsics[:, :2, -1] <= 1) ), "Intrinsics should be expressed in resolution-independent normalized image coordinates." c2w_src = torch.linalg.inv(extrinsics_src_flat) # transform coordinates from the source camera's coordinate system to the coordinate system of the respective camera extrinsics_rel = torch.einsum( "vnm,vmp->vnp", extrinsics_flat, c2w_src ) intrinsics[:, :2] *= extrinsics_flat.new_tensor( [ target_size[1], # w target_size[0], # h ] ).view(1, -1, 1) centers, rays, grid_cam = get_center_and_ray( img_h=target_size[0], img_w=target_size[1], pose=extrinsics_rel[:, :3, :], intr=intrinsics, ) rays = torch.nn.functional.normalize(rays, dim=-1) plucker = torch.cat((rays, torch.cross(centers, rays, dim=-1)), dim=-1) plucker = plucker.permute(0, 2, 1).reshape(plucker.shape[0], -1, *target_size) # Reshape back to batch dimension if needed if has_batch_dim: plucker = plucker.reshape(batch_size, num_cameras, *plucker.shape[1:]) return plucker def get_value_dict( curr_imgs, curr_imgs_clip, curr_input_frame_indices, curr_c2ws, curr_Ks, curr_input_camera_indices, all_c2ws, camera_scale, ): assert sorted(curr_input_camera_indices) == sorted( range(len(curr_input_camera_indices)) ) H, W, T, F = curr_imgs.shape[-2], curr_imgs.shape[-1], len(curr_imgs), 8 value_dict = {} value_dict["cond_frames_without_noise"] = curr_imgs_clip[curr_input_frame_indices] value_dict["cond_frames"] = curr_imgs + 0.0 * torch.randn_like(curr_imgs) value_dict["cond_frames_mask"] = torch.zeros(T, dtype=torch.bool) value_dict["cond_frames_mask"][curr_input_frame_indices] = True value_dict["cond_aug"] = 0.0 if curr_c2ws.shape[-1] == 3: c2w = to_hom_pose(curr_c2ws.float()) else: c2w = curr_c2ws w2c = torch.linalg.inv(c2w) # camera centering ref_c2ws = all_c2ws camera_dist_2med = torch.norm( ref_c2ws[:, :3, 3] - ref_c2ws[:, :3, 3].median(0, keepdim=True).values, dim=-1, ) valid_mask = camera_dist_2med <= torch.clamp( torch.quantile(camera_dist_2med, 0.97) * 10, max=1e6, ) c2w[:, :3, 3] -= ref_c2ws[valid_mask, :3, 3].mean(0, keepdim=True) w2c = torch.linalg.inv(c2w) # camera normalization camera_dists = c2w[:, :3, 3].clone() translation_scaling_factor = ( camera_scale if torch.isclose( torch.norm(camera_dists[0]), torch.zeros(1), atol=1e-5, ).any() else (camera_scale / torch.norm(camera_dists[0])) ) w2c[:, :3, 3] *= translation_scaling_factor c2w[:, :3, 3] *= translation_scaling_factor value_dict["plucker_coordinate"] = get_plucker_coordinates( extrinsics_src=w2c[0], extrinsics=w2c, intrinsics=curr_Ks.float().clone(), target_size=(H // F, W // F), ) value_dict["c2w"] = c2w value_dict["K"] = curr_Ks value_dict["camera_mask"] = torch.zeros(T, dtype=torch.bool) value_dict["camera_mask"][curr_input_camera_indices] = True return value_dict def parse_meta_data(file_path, image_height=288, image_width=512): with open(file_path, 'r') as file: lines = file.readlines() # First line is the video URL video_url = lines[0].strip() line = lines[1] data = line.strip().split() # Construct the camera intrinsics matrix K focal_length_x = float(data[1]) focal_length_y = float(data[2]) principal_point_x = float(data[3]) principal_point_y = float(data[4]) original_K = [ [focal_length_x, 0, principal_point_x], [0, focal_length_y, principal_point_y], [0, 0, 1] ] K = [ [focal_length_x * image_width, 0, principal_point_x * image_width], [0, focal_length_y * image_height, principal_point_y * image_height], [0, 0, 1] ] # Initialize a list to store frame data timestamp_to_c2ws = {} timestamps = [] # Process each frame line for line in lines[1:]: data = line.strip().split() timestamp = int(data[0]) R_t = [float(x) for x in data[7:]] P = [ R_t[0:4], R_t[4:8], R_t[8:12], [0, 0, 0, 1] ] timestamp_to_c2ws[timestamp] = np.array(P) timestamps.append(timestamp) return timestamps, np.array(K), timestamp_to_c2ws, original_K def get_wh_with_fixed_shortest_side(w, h, size): # size is smaller or equal to zero, we return original w h if size is None or size <= 0: return w, h if w < h: new_w = size new_h = int(size * h / w) else: new_h = size new_w = int(size * w / h) return new_w, new_h def get_resizing_factor( target_shape: Tuple[int, int], # H, W current_shape: Tuple[int, int], # H, W cover_target: bool = True, # If True, the output shape will fully cover the target shape. # If No, the target shape will fully cover the output shape. ) -> float: r_bound = target_shape[1] / target_shape[0] aspect_r = current_shape[1] / current_shape[0] if r_bound >= 1.0: if cover_target: if aspect_r >= r_bound: factor = min(target_shape) / min(current_shape) elif aspect_r < 1.0: factor = max(target_shape) / min(current_shape) else: factor = max(target_shape) / max(current_shape) else: if aspect_r >= r_bound: factor = max(target_shape) / max(current_shape) elif aspect_r < 1.0: factor = min(target_shape) / max(current_shape) else: factor = min(target_shape) / min(current_shape) else: if cover_target: if aspect_r <= r_bound: factor = min(target_shape) / min(current_shape) elif aspect_r > 1.0: factor = max(target_shape) / min(current_shape) else: factor = max(target_shape) / max(current_shape) else: if aspect_r <= r_bound: factor = max(target_shape) / max(current_shape) elif aspect_r > 1.0: factor = min(target_shape) / max(current_shape) else: factor = min(target_shape) / min(current_shape) return factor def transform_img_and_K( image: torch.Tensor, size: Union[int, Tuple[int, int]], scale: float = 1.0, center: Tuple[float, float] = (0.5, 0.5), K: Union[torch.Tensor, np.ndarray, None] = None, size_stride: int = 1, mode: str = "crop", ): assert mode in [ "crop", "pad", "stretch", ], f"mode should be one of ['crop', 'pad', 'stretch'], got {mode}" h, w = image.shape[-2:] if isinstance(size, (tuple, list)): # => if size is a tuple or list, we first rescale to fully cover the `size` # area and then crop the `size` area from the rescale image W, H = size else: # => if size is int, we rescale the image to fit the shortest side to size # => if size is None, no rescaling is applied W, H = get_wh_with_fixed_shortest_side(w, h, size) W, H = ( math.floor(W / size_stride + 0.5) * size_stride, math.floor(H / size_stride + 0.5) * size_stride, ) if mode == "stretch": rh, rw = H, W else: rfs = get_resizing_factor( (H, W), (h, w), cover_target=mode != "pad", ) (rh, rw) = [int(np.ceil(rfs * s)) for s in (h, w)] rh, rw = int(rh / scale), int(rw / scale) image = torch.nn.functional.interpolate( image, (rh, rw), mode="area", antialias=False ) cy_center = int(center[1] * image.shape[-2]) cx_center = int(center[0] * image.shape[-1]) if mode != "pad": ct = max(0, cy_center - H // 2) cl = max(0, cx_center - W // 2) ct = min(ct, image.shape[-2] - H) cl = min(cl, image.shape[-1] - W) image = TF.crop(image, top=ct, left=cl, height=H, width=W) pl, pt = 0, 0 else: pt = max(0, H // 2 - cy_center) pl = max(0, W // 2 - cx_center) pb = max(0, H - pt - image.shape[-2]) pr = max(0, W - pl - image.shape[-1]) image = TF.pad( image, [pl, pt, pr, pb], ) cl, ct = 0, 0 if K is not None: K = K.clone() # K[:, :2, 2] += K.new_tensor([pl, pt]) if torch.all(K[:, :2, -1] >= 0) and torch.all(K[:, :2, -1] <= 1): K[:, :2] *= K.new_tensor([rw, rh])[None, :, None] # normalized K else: K[:, :2] *= K.new_tensor([rw / w, rh / h])[None, :, None] # unnormalized K K[:, :2, 2] += K.new_tensor([pl - cl, pt - ct]) return image, K def load_img_and_K( image_path_or_size: Union[str, torch.Size], size: Optional[Union[int, Tuple[int, int]]], scale: float = 1.0, center: Tuple[float, float] = (0.5, 0.5), K: Union[torch.Tensor, np.ndarray, None] = None, size_stride: int = 1, center_crop: bool = False, image_as_tensor: bool = True, context_rgb: Union[np.ndarray, None] = None, device: str = "cuda", ): if isinstance(image_path_or_size, torch.Size): image = Image.new("RGBA", image_path_or_size[::-1]) else: image = Image.open(image_path_or_size).convert("RGBA") w, h = image.size if size is None: size = (w, h) image = np.array(image).astype(np.float32) / 255 if image.shape[-1] == 4: rgb, alpha = image[:, :, :3], image[:, :, 3:] if context_rgb is not None: image = rgb * alpha + context_rgb * (1 - alpha) else: image = rgb * alpha + (1 - alpha) image = image.transpose(2, 0, 1) image = torch.from_numpy(image).to(dtype=torch.float32) image = image.unsqueeze(0) if isinstance(size, (tuple, list)): # => if size is a tuple or list, we first rescale to fully cover the `size` # area and then crop the `size` area from the rescale image W, H = size else: # => if size is int, we rescale the image to fit the shortest side to size # => if size is None, no rescaling is applied W, H = get_wh_with_fixed_shortest_side(w, h, size) W, H = ( math.floor(W / size_stride + 0.5) * size_stride, math.floor(H / size_stride + 0.5) * size_stride, ) rfs = get_resizing_factor((math.floor(H * scale), math.floor(W * scale)), (h, w)) resize_size = rh, rw = [int(np.ceil(rfs * s)) for s in (h, w)] image = torch.nn.functional.interpolate( image, resize_size, mode="area", antialias=False ) if scale < 1.0: pw = math.ceil((W - resize_size[1]) * 0.5) ph = math.ceil((H - resize_size[0]) * 0.5) image = F.pad(image, (pw, pw, ph, ph), "constant", 1.0) cy_center = int(center[1] * image.shape[-2]) cx_center = int(center[0] * image.shape[-1]) if center_crop: side = min(H, W) ct = max(0, cy_center - side // 2) cl = max(0, cx_center - side // 2) ct = min(ct, image.shape[-2] - side) cl = min(cl, image.shape[-1] - side) image = TF.crop(image, top=ct, left=cl, height=side, width=side) else: ct = max(0, cy_center - H // 2) cl = max(0, cx_center - W // 2) ct = min(ct, image.shape[-2] - H) cl = min(cl, image.shape[-1] - W) image = TF.crop(image, top=ct, left=cl, height=H, width=W) if K is not None: K = K.clone() if torch.all(K[:2, -1] >= 0) and torch.all(K[:2, -1] <= 1): K[:2] *= K.new_tensor([rw, rh])[:, None] # normalized K else: K[:2] *= K.new_tensor([rw / w, rh / h])[:, None] # unnormalized K K[:2, 2] -= K.new_tensor([cl, ct]) if image_as_tensor: # tensor of shape (1, 3, H, W) with values ranging from (-1, 1) image = image.to(device) * 2.0 - 1.0 else: # PIL Image with values ranging from (0, 255) image = image.permute(0, 2, 3, 1).numpy()[0] image = Image.fromarray((image * 255).astype(np.uint8)) return image, K def geodesic_distance(extrinsic1: Union[np.ndarray, torch.Tensor], extrinsic2: Union[np.ndarray, torch.Tensor], weight_translation: float = 0.01,): """ Computes the geodesic distance between two camera poses in SE(3). Parameters: extrinsic1 (Union[np.ndarray, torch.Tensor]): 4x4 extrinsic matrix of the first pose. extrinsic2 (Union[np.ndarray, torch.Tensor]): 4x4 extrinsic matrix of the second pose. Returns: Union[float, torch.Tensor]: Geodesic distance between the two poses. """ if torch.is_tensor(extrinsic1): # Extract the rotation and translation components R1 = extrinsic1[:3, :3] t1 = extrinsic1[:3, 3] R2 = extrinsic2[:3, :3] t2 = extrinsic2[:3, 3] # Compute the translation distance (Euclidean distance) translation_distance = torch.norm(t1 - t2) # Compute the relative rotation matrix R_relative = torch.matmul(R1.T, R2) # Compute the angular distance from the trace of the relative rotation matrix trace_value = torch.trace(R_relative) # Clamp the trace value to avoid numerical issues trace_value = torch.clamp(trace_value, -1.0, 3.0) angular_distance = torch.acos((trace_value - 1) / 2) else: # Extract the rotation and translation components R1 = extrinsic1[:3, :3] t1 = extrinsic1[:3, 3] R2 = extrinsic2[:3, :3] t2 = extrinsic2[:3, 3] # Compute the translation distance (Euclidean distance) translation_distance = np.linalg.norm(t1 - t2) # Compute the relative rotation matrix R_relative = np.dot(R1.T, R2) # Compute the angular distance from the trace of the relative rotation matrix trace_value = np.trace(R_relative) # Clamp the trace value to avoid numerical issues trace_value = np.clip(trace_value, -1.0, 3.0) angular_distance = np.arccos((trace_value - 1) / 2) # Combine the two distances geodesic_dist = translation_distance*weight_translation + angular_distance return geodesic_dist def inverse_geodesic_distance(extrinsic1, extrinsic2, weight_translation=0.01): """ Computes the inverse geodesic distance between two camera poses in SE(3). Parameters: extrinsic1 (np.ndarray): 4x4 extrinsic matrix of the first pose. extrinsic2 (np.ndarray): 4x4 extrinsic matrix of the second pose. Returns: float: Inverse geodesic distance between the two poses. """ # Compute the geodesic distance geodesic_dist = geodesic_distance(extrinsic1, extrinsic2, weight_translation) # Compute the inverse geodesic distance inverse_geodesic_dist = 1.0 / (geodesic_dist + 1e-6) return inverse_geodesic_dist def average_camera_pose(camera_poses): """ Compute a better average of camera poses in SE(3). Args: camera_poses: List or array of camera poses, each a 4x4 matrix Returns: Average camera pose as a 4x4 matrix """ # Extract rotation and translation components rotations = camera_poses[:, :3, :3].detach().cpu().numpy() translations = camera_poses[:, :3, 3].detach().cpu().numpy() # Average translation with simple mean avg_translation = np.mean(translations, axis=0) # Convert rotations to quaternions for better averaging import scipy.spatial.transform as transform quats = [transform.Rotation.from_matrix(R).as_quat() for R in rotations] # Ensure quaternions are in the same hemisphere to avoid issues with averaging for i in range(1, len(quats)): if np.dot(quats[0], quats[i]) < 0: quats[i] = -quats[i] # Average the quaternions and convert back to rotation matrix avg_quat = np.mean(quats, axis=0) avg_quat = avg_quat / np.linalg.norm(avg_quat) # Normalize avg_rotation = transform.Rotation.from_quat(avg_quat).as_matrix() # Construct the average pose avg_pose = np.eye(4) avg_pose[:3, :3] = avg_rotation avg_pose[:3, 3] = avg_translation return avg_pose def encode_image( image, image_encoder, device, dtype, ) -> torch.Tensor: image = image.to(device=device, dtype=dtype) image_embeddings = image_encoder(image) return image_embeddings def encode_vae_image( image, vae, device, dtype, ): image = image.to(device=device, dtype=dtype) image_latents = vae.encode(image, 1) return image_latents def do_sample( model, ae, denoiser, sampler, c, uc, c2w, K, cond_frames_mask, H=576, W=768, C=4, F=8, T=8, cfg=2.0, decoding_t=1, verbose=True, global_pbar=None, return_latents=False, device: str = "cuda", **_, ): num_samples = [1, T] with torch.inference_mode(), torch.autocast("cuda"): additional_model_inputs = {"num_frames": T} additional_sampler_inputs = { "c2w": c2w.to("cuda"), "K": K.to("cuda"), "input_frame_mask": cond_frames_mask.to("cuda"), } if global_pbar is not None: additional_sampler_inputs["global_pbar"] = global_pbar shape = (math.prod(num_samples), C, H // F, W // F) randn = torch.randn(shape).to(device) samples_z = sampler( lambda input, sigma, c: denoiser( model, input, sigma, c, **additional_model_inputs, ), randn, scale=cfg, cond=c, uc=uc, verbose=verbose, **additional_sampler_inputs, ) if samples_z is None: return samples = ae.decode(samples_z, decoding_t) if return_latents: return samples, samples_z return samples def decode_output( samples, T, indices=None, ): # decode model output into dict if it is not if isinstance(samples, dict): # model with postprocessor and outputs dict q`` for sample, value in samples.items(): if isinstance(value, torch.Tensor): value = value.detach().cpu() elif isinstance(value, np.ndarray): value = torch.from_numpy(value) else: value = torch.tensor(value) if indices is not None and value.shape[0] == T: value = value[indices] samples[sample] = value else: # model without postprocessor and outputs tensor (rgb) samples = samples.detach().cpu() if indices is not None and samples.shape[0] == T: samples = samples[indices] samples = {"samples-rgb/image": samples} return samples def select_frames(timestamps, min_num_frames=2, skip_frame=10, random_start=False): """ Select frames from a video sequence based on defined criteria. Args: timestamps: List of timestamps for the frames min_num_frames: Minimum number of frames required skip_frame: Number of frames to skip between selections random_start: If True, start from a random frame Returns: tuple: (selected_frame_indices, selected_frame_timestamps) or (None, None) if criteria not met """ num_frames = len(timestamps) if num_frames < min_num_frames: print(f"[Worker PID={os.getpid()}] Episode has less than {min_num_frames} frames") return None, None # Decide on start/end frames if num_frames < 2: print(f"[Worker PID={os.getpid()}] Episode has less than 2 frames") return None, None elif num_frames < skip_frame: cur_skip_frame = num_frames - 1 else: cur_skip_frame = skip_frame if random_start: start_frame = np.random.randint(0, skip_frame) else: start_frame = 0 # Gather frame indices selected_frame_indices = list(range(start_frame, num_frames, cur_skip_frame)) selected_frame_timestamps = [timestamps[i] for i in selected_frame_indices] return selected_frame_indices, selected_frame_timestamps def tensor2im(input_image, imtype=np.uint8): if not isinstance(input_image, np.ndarray): if isinstance(input_image, torch.Tensor): # get the data from a variable image_tensor = input_image.data else: return input_image image_numpy = image_tensor[0].clamp(0.0, 1.0).cpu().float().numpy() # convert it into a numpy array image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0 # post-processing: tranpose and scaling else: # if it is a numpy array, do nothing image_numpy = input_image return image_numpy.astype(imtype) class LatentStorer: def __init__(self): self.latent = None def __call__(self, i, t, latent): self.latent = latent def sobel_filter(disp, mode="sobel", beta=10.0): sobel_grad = kornia.filters.spatial_gradient(disp, mode=mode, normalized=False) sobel_mag = torch.sqrt(sobel_grad[:, :, 0, Ellipsis] ** 2 + sobel_grad[:, :, 1, Ellipsis] ** 2) alpha = torch.exp(-1.0 * beta * sobel_mag).detach() return alpha def apply_colormap(image, cmap="viridis"): colormap = cm.get_cmap(cmap) colormap = torch.tensor(colormap.colors).to(image.device) image_long = (image * 255).long() image_long_min = torch.min(image_long) image_long_max = torch.max(image_long) assert image_long_min >= 0, f"the min value is {image_long_min}" assert image_long_max <= 255, f"the max value is {image_long_max}" return colormap[image_long[..., 0]] def apply_depth_colormap( depth, near_plane=None, far_plane=None, cmap="viridis", ): near_plane = near_plane or float(torch.min(depth)) far_plane = far_plane or float(torch.max(depth)) depth = (depth - near_plane) / (far_plane - near_plane + 1e-10) depth = torch.clip(depth, 0, 1) colored_image = apply_colormap(depth, cmap=cmap) return colored_image def save_video(video, path, fps=10): video = video.permute(0, 2, 3, 1) video_codec = "libx264" video_options = { "crf": "23", # Constant Rate Factor (lower value = higher quality, 18 is a good balance) "preset": "slow", } write_video(str(path), video, fps=fps, video_codec=video_codec, options=video_options) def visualize_camera_poses(camera_poses, axis_length=0.1): """ Visualizes a set of camera poses in 3D using Matplotlib. Parameters ---------- camera_poses : np.ndarray An array of shape (N, 4, 4) containing N camera poses. Each pose is a 4x4 transformation matrix. axis_length : float Length of the camera axes to draw. """ if isinstance(camera_poses, torch.Tensor): camera_poses = camera_poses.detach().cpu().numpy() # Create a 3D figure fig = plt.figure() ax = fig.add_subplot(111, projection='3d') # Iterate over all camera poses for i in range(camera_poses.shape[0]): # Extract rotation (R) and translation (t) R = camera_poses[i][:3, :3] t = camera_poses[i][:3, 3] # Plot the camera center ax.scatter(t[0], t[1], t[2], c='k', marker='o', s=20) # Define the end-points of each local axis x_axis_end = t + R[:, 0] * axis_length y_axis_end = t + R[:, 1] * axis_length z_axis_end = t + R[:, 2] * axis_length # Draw the axes as lines ax.plot([t[0], x_axis_end[0]], [t[1], x_axis_end[1]], [t[2], x_axis_end[2]], color='r') # X-axis (red) ax.plot([t[0], y_axis_end[0]], [t[1], y_axis_end[1]], [t[2], y_axis_end[2]], color='g') # Y-axis (green) ax.plot([t[0], z_axis_end[0]], [t[1], z_axis_end[1]], [t[2], z_axis_end[2]], color='b') # Z-axis (blue) # Make axes have equal scale set_axes_equal(ax) ax.set_title("Camera Poses Visualization") ax.set_xlabel("X") ax.set_ylabel("Y") ax.set_zlabel("Z") plt.show() def set_axes_equal(ax): """ Make axes of 3D plot have equal scale so that spheres appear as spheres, cubes as cubes, etc. This is a workaround to Matplotlib's set_aspect('equal') which is not supported in 3D. """ x_limits = ax.get_xlim3d() y_limits = ax.get_ylim3d() z_limits = ax.get_zlim3d() x_range = x_limits[1] - x_limits[0] y_range = y_limits[1] - y_limits[0] z_range = z_limits[1] - z_limits[0] max_range = max(x_range, y_range, z_range) x_middle = np.mean(x_limits) y_middle = np.mean(y_limits) z_middle = np.mean(z_limits) ax.set_xlim3d([x_middle - 0.5 * max_range, x_middle + 0.5 * max_range]) ax.set_ylim3d([y_middle - 0.5 * max_range, y_middle + 0.5 * max_range]) ax.set_zlim3d([z_middle - 0.5 * max_range, z_middle + 0.5 * max_range]) def tensor_to_pil(image): if isinstance(image, torch.Tensor): if image.dim() == 4: image = image.squeeze(0) image = image.permute(1, 2, 0).detach().cpu().numpy() # Detect the range of the input tensor if image.min() < -0.1: # If we have negative values, assume [-1, 1] range image = (image + 1) / 2.0 # Convert from [-1, 1] to [0, 1] # Otherwise, assume it's already in [0, 1] range image = (image * 255) image = np.clip(image, 0, 255) image = image.astype(np.uint8) return Image.fromarray(image) def center_crop_pil_image(input_image, target_width=1024, target_height=576): w, h = input_image.size h_ratio = h / target_height w_ratio = w / target_width if h_ratio > w_ratio: h = int(h / w_ratio) if h < target_height: h = target_height input_image = input_image.resize((target_width, h), Image.Resampling.LANCZOS) else: w = int(w / h_ratio) if w < target_width: w = target_width input_image = input_image.resize((w, target_height), Image.Resampling.LANCZOS) return ImageOps.fit(input_image, (target_width, target_height), Image.BICUBIC) def resize_pil_image(img, long_edge_size): S = max(img.size) if S > long_edge_size: interp = PIL.Image.LANCZOS elif S <= long_edge_size: interp = PIL.Image.BICUBIC new_size = tuple(int(round(x*long_edge_size/S)) for x in img.size) return img.resize(new_size, interp) def visualize_surfels( surfels, draw_normals=False, normal_scale=20, disk_resolution=16, disk_alpha=0.5 ): """ Visualize surfels as 2D disks oriented by their normals in 3D using matplotlib. Args: surfels (list of Surfel): Each Surfel has at least: - position: (x, y, z) - normal: (nx, ny, nz) - radius: scalar - color: (R, G, B) in [0..255] (optional) draw_normals (bool): If True, draws the surfel normals as quiver arrows. normal_scale (float): Scale factor for the normal arrows. disk_resolution (int): Number of segments to approximate each disk. disk_alpha (float): Alpha (transparency) for the filled disks. """ fig = plt.figure() ax = fig.add_subplot(111, projection='3d') # Prepare arrays for optional quiver (if draw_normals=True) positions = [] normals = [] # We'll accumulate 3D polygons in a list for Poly3DCollection polygons = [] polygon_colors = [] for s in surfels: # --- Extract surfel data --- position = s.position normal = s.normal radius = s.radius if isinstance(position, torch.Tensor): x, y, z = position.detach().cpu().numpy() nx, ny, nz = normal.detach().cpu().numpy() radius = radius.detach().cpu().numpy() else: x, y, z = position nx, ny, nz = normal radius = radius # Convert color from [0..255] to [0..1], or use default if s.color is None: color = (0.2, 0.6, 1.0) # Light blue else: r, g, b = s.color color = (r/255.0, g/255.0, b/255.0) # --- Build local coordinate axes for the disk --- normal = np.array([nx, ny, nz], dtype=float) norm_len = np.linalg.norm(normal) # Skip degenerate normals to avoid nan if norm_len < 1e-12: continue normal /= norm_len # Pick an 'up' vector that is not too close to the normal # so we can build a tangent plane up = np.array([0, 0, 1], dtype=float) if abs(normal.dot(up)) > 0.9: up = np.array([0, 1, 0], dtype=float) # xAxis = normal x up xAxis = np.cross(normal, up) xAxis /= np.linalg.norm(xAxis) # yAxis = normal x xAxis yAxis = np.cross(normal, xAxis) yAxis /= np.linalg.norm(yAxis) # --- Create a circle of 'disk_resolution' segments in local 2D coords --- angles = np.linspace(0, 2*np.pi, disk_resolution, endpoint=False) circle_points_3d = [] for theta in angles: # local 2D circle: (r*cosθ, r*sinθ) px = radius * np.cos(theta) py = radius * np.sin(theta) # transform to 3D world space: position + px*xAxis + py*yAxis world_pt = np.array([x, y, z]) + px * xAxis + py * yAxis circle_points_3d.append(world_pt) # We have a list of [x, y, z]. For a filled polygon, Poly3DCollection # wants them as a single Nx3 array. circle_points_3d = np.array(circle_points_3d) polygons.append(circle_points_3d) polygon_colors.append(color) # Collect positions and normals for quiver (if used) positions.append([x, y, z]) normals.append(normal) # --- Draw the disks as polygons --- poly_collection = Poly3DCollection( polygons, facecolors=polygon_colors, edgecolors='k', # black edge linewidths=0.5, alpha=disk_alpha ) ax.add_collection3d(poly_collection) # --- Optionally draw normal vectors (quiver) --- if draw_normals and len(positions) > 0: X = [p[0] for p in positions] Y = [p[1] for p in positions] Z = [p[2] for p in positions] Nx = [n[0] for n in normals] Ny = [n[1] for n in normals] Nz = [n[2] for n in normals] # Note: If your scene is large, you may want to increase `length`. ax.quiver( X, Y, Z, Nx, Ny, Nz, length=normal_scale, color='red', normalize=True ) # --- Axis labels, aspect ratio, etc. --- ax.set_xlabel('X') ax.set_ylabel('Y') ax.set_zlabel('Z') try: ax.set_box_aspect((1, 1, 1)) except AttributeError: pass # older MPL versions plt.title("Surfels as Disks (Oriented by Normal)") plt.show() def visualize_pointcloud( points, colors=None, title='Point Cloud', point_size=1, alpha=1.0 ): """ Visualize a 3D point cloud using Matplotlib, with an option to provide per-point RGB or RGBA colors, ensuring equal scaling for the x, y, and z axes. Parameters ---------- points : np.ndarray or torch.Tensor A numpy array (or Tensor) of shape [N, 3] where each row is a 3D point (x, y, z). colors : None, str, or np.ndarray - If None, a default single color ('blue') is used. - If a string, that color will be used for all points. - If a numpy array, it should have shape [N, 3] or [N, 4], where each row corresponds to the color of the matching point in `points`. Values should be in the range [0, 1] if using floats. title : str, optional The title of the plot. Default is 'Point Cloud'. point_size : float, optional The size of the points in the scatter plot. Default is 1. alpha : float, optional The overall alpha (transparency) value for the points. Default is 1.0. Examples -------- >>> import numpy as np >>> # Generate random points >>> pts = np.random.rand(1000, 3) >>> # Generate random colors in [0,1] >>> cols = np.random.rand(1000, 3) >>> visualize_pointcloud(pts, colors=cols, title="Random Point Cloud with Colors") """ # Convert Torch tensors to NumPy arrays if needed if isinstance(points, torch.Tensor): points = points.detach().cpu().numpy() if isinstance(colors, torch.Tensor): colors = colors.detach().cpu().numpy() # Flatten points if they are in a higher-dimensional array if len(points.shape) > 2: points = points.reshape(-1, 3) if colors is not None and isinstance(colors, np.ndarray) and len(colors.shape) > 2: colors = colors.reshape(-1, colors.shape[-1]) # Validate shape of points if points.shape[1] != 3: raise ValueError("`points` array must have shape [N, 3].") # Validate or set colors if colors is None: colors = 'blue' elif isinstance(colors, np.ndarray): colors = np.asarray(colors) if colors.shape[0] != points.shape[0]: raise ValueError( "Colors array length must match the number of points." ) if colors.shape[1] not in [3, 4]: raise ValueError( "Colors array must have shape [N, 3] or [N, 4]." ) # Extract coordinates x = points[:, 0] y = points[:, 1] z = points[:, 2] # Create a 3D figure fig = plt.figure(figsize=(8, 6)) ax = fig.add_subplot(111, projection='3d') # Scatter plot with specified or per-point colors ax.scatter(x, y, z, c=colors, s=point_size, alpha=alpha) # Set labels and title ax.set_xlabel('X') ax.set_ylabel('Y') ax.set_zlabel('Z') ax.set_title(title) # Ensure all axes have the same scale max_range = np.array([x.max() - x.min(), y.max() - y.min(), z.max() - z.min()]).max() / 2.0 mid_x = (x.max() + x.min()) * 0.5 mid_y = (y.max() + y.min()) * 0.5 mid_z = (z.max() + z.min()) * 0.5 ax.set_xlim(mid_x - max_range, mid_x + max_range) ax.set_ylim(mid_y - max_range, mid_y + max_range) ax.set_zlim(mid_z - max_range, mid_z + max_range) # Adjust viewing angle for better visibility ax.view_init(elev=20., azim=30) plt.tight_layout() plt.show() def visualize_depth(depth_image, file_name="rendered_depth.png", visualization_dir="visualization", size=(512, 288)): """ Visualize a depth map as a grayscale image. Parameters ---------- depth_image : np.ndarray A 2D array of depth values. visualization_dir : str The directory to save the visualization image. Returns ------- PIL.Image The visualization image. """ # Normalize the depth values for visualization depth_min = depth_image.min() depth_max = depth_image.max() print(f"Depth min: {depth_min}, max: {depth_max}") depth_image = np.clip(depth_image, 0, depth_max) depth_vis = (depth_image - depth_min) / (depth_max - depth_min) depth_vis = (depth_vis * 255).astype(np.uint8) # Convert the depth image to a PIL image depth_vis_img = Image.fromarray(depth_vis, mode='L') depth_vis_img = depth_vis_img.resize(size, Image.NEAREST) # Save the visualization image depth_vis_img.save(os.path.join(visualization_dir, file_name)) return depth_vis_img class Surfel: def __init__(self, position, normal, radius=1.0, color=None): """ position: (x, y, z) normal: (nx, ny, nz) radius: scalar color: (r, g, b) or None """ self.position = position self.normal = normal self.radius = radius self.color = color def __repr__(self): return (f"Surfel(position={self.position}, " f"normal={self.normal}, radius={self.radius}, " f"color={self.color})") class Octree: def __init__(self, points, indices=None, bbox=None, max_points=10): self.points = points if indices is None: indices = np.arange(points.shape[0]) self.indices = indices if bbox is None: min_bound = points.min(axis=0) max_bound = points.max(axis=0) center = (min_bound + max_bound) / 2 half_size = np.max(max_bound - min_bound) / 2 bbox = (center, half_size) self.center, self.half_size = bbox self.children = [] # 存储子节点 self.max_points = max_points if len(self.indices) > self.max_points: self.subdivide() def subdivide(self): cx, cy, cz = self.center hs = self.half_size / 2 offsets = np.array([[dx, dy, dz] for dx in (-hs, hs) for dy in (-hs, hs) for dz in (-hs, hs)]) for offset in offsets: child_center = self.center + offset child_indices = [] for idx in self.indices: p = self.points[idx] if np.all(np.abs(p - child_center) <= hs): child_indices.append(idx) child_indices = np.array(child_indices) if len(child_indices) > 0: child = Octree(self.points, indices=child_indices, bbox=(child_center, hs), max_points=self.max_points) self.children.append(child) self.indices = None def sphere_intersects_node(self, center, r): diff = np.abs(center - self.center) max_diff = diff - self.half_size max_diff = np.maximum(max_diff, 0) dist_sq = np.sum(max_diff**2) return dist_sq <= r*r def query_ball_point(self, point, r): results = [] if not self.sphere_intersects_node(point, r): return results if len(self.children) == 0: if self.indices is not None: for idx in self.indices: if np.linalg.norm(self.points[idx] - point) <= r: results.append(idx) return results else: for child in self.children: results.extend(child.query_ball_point(point, r)) return results