import sys import os import numpy as np import scipy.interpolate import PIL import torch import matplotlib.pyplot as plt import moviepy.editor as mpy sys.path.append('submodules/mast3r/dust3r') from dust3r.utils.image import heif_support_enabled, exif_transpose, _resize_pil_image, ImgNorm from dust3r.image_pairs import make_pairs from dust3r.inference import inference from dust3r.cloud_opt import global_aligner, GlobalAlignerMode sys.path.append('.') from src.utils.cuda_splatting import render, DummyPipeline from src.utils.gaussian_model import GaussianModel from src.utils.camera_utils import get_scaled_camera from src.losses import merge_and_split_predictions from src.utils.camera_utils import move_c2w_along_z from einops import rearrange LABELS = ['wall', 'floor', 'ceiling', 'chair', 'table', 'sofa', 'bed', 'other'] NUM_LABELS = len(LABELS) + 1 PALLETE = plt.get_cmap('tab10', NUM_LABELS) COLORS_LIST = [PALLETE(i)[:3] for i in range(NUM_LABELS)] COLORS = torch.tensor(COLORS_LIST, dtype=torch.float32) def load_images(folder_or_list, size, square_ok=False, verbose=True, save_dir=None): """ open and convert all images in a list or folder to proper input format for DUSt3R """ if isinstance(folder_or_list, str): if verbose: print(f'>> Loading images from {folder_or_list}') root, folder_content = folder_or_list, sorted(os.listdir(folder_or_list)) elif isinstance(folder_or_list, list): if verbose: print(f'>> Loading a list of {len(folder_or_list)} images') root, folder_content = '', folder_or_list else: raise ValueError(f'bad {folder_or_list=} ({type(folder_or_list)})') supported_images_extensions = ['.jpg', '.jpeg', '.png'] if heif_support_enabled: supported_images_extensions += ['.heic', '.heif'] supported_images_extensions = tuple(supported_images_extensions) imgs = [] for path in folder_content: if not path.lower().endswith(supported_images_extensions): continue img = exif_transpose(PIL.Image.open(os.path.join(root, path))).convert('RGB') W1, H1 = img.size if size == 224: # resize short side to 224 (then crop) img = _resize_pil_image(img, round(size * max(W1/H1, H1/W1))) else: # resize long side to 512 img = _resize_pil_image(img, size) W, H = img.size cx, cy = W//2, H//2 if size == 224: half = min(cx, cy) img = img.crop((cx-half, cy-half, cx+half, cy+half)) else: halfw, halfh = ((2*cx)//32)*16, ((2*cy)//32)*16 if not (square_ok) and W == H: halfh = 3*halfw/4 img = img.crop((cx-halfw, cy-halfh, cx+halfw, cy+halfh)) W2, H2 = img.size if verbose: print(f' - adding {path} with resolution {W1}x{H1} --> {W2}x{H2}') # Save the processed image if save_dir is provided if save_dir: os.makedirs(save_dir, exist_ok=True) save_path = os.path.join(save_dir, f"processed_{len(imgs):03d}.png") img.save(save_path) if verbose: print(f' - saved processed image to {save_path}') imgs.append(dict(img=ImgNorm(img)[None], true_shape=np.int32( [img.size[::-1]]), idx=len(imgs), instance=str(len(imgs)))) assert imgs, 'no images foud at '+root if verbose: print(f' (Found {len(imgs)} images)') return imgs def normalize(x): """Normalization helper function.""" return x / np.linalg.norm(x) def viewmatrix(lookdir, up, position): """Construct lookat view matrix.""" vec2 = normalize(lookdir) vec0 = normalize(np.cross(up, vec2)) vec1 = normalize(np.cross(vec2, vec0)) m = np.stack([vec0, vec1, vec2, position], axis=1) return m def poses_to_points(poses, dist): """Converts from pose matrices to (position, lookat, up) format.""" pos = poses[:, :3, -1] lookat = poses[:, :3, -1] - dist * poses[:, :3, 2] up = poses[:, :3, -1] + dist * poses[:, :3, 1] return np.stack([pos, lookat, up], 1) def points_to_poses(points): """Converts from (position, lookat, up) format to pose matrices.""" return np.array([viewmatrix(p - l, u - p, p) for p, l, u in points]) def interp(points, n, k, s): """Runs multidimensional B-spline interpolation on the input points.""" sh = points.shape pts = np.reshape(points, (sh[0], -1)) k = min(k, sh[0] - 1) tck, _ = scipy.interpolate.splprep(pts.T, k=k, s=s) u = np.linspace(0, 1, n, endpoint=False) new_points = np.array(scipy.interpolate.splev(u, tck)) new_points = np.reshape(new_points.T, (n, sh[1], sh[2])) return new_points def generate_interpolated_path(poses, n_interp, spline_degree=5, smoothness=.03, rot_weight=.1): """Creates a smooth spline path between input keyframe camera poses. Spline is calculated with poses in format (position, lookat-point, up-point). Args: poses: (n, 3, 4) array of input pose keyframes. n_interp: returned path will have n_interp * (n - 1) total poses. spline_degree: polynomial degree of B-spline. smoothness: parameter for spline smoothing, 0 forces exact interpolation. rot_weight: relative weighting of rotation/translation in spline solve. Returns: Array of new camera poses with shape (n_interp * (n - 1), 3, 4). """ points = poses_to_points(poses, dist=rot_weight) new_points = interp(points, n_interp * (points.shape[0] - 1), k=spline_degree, s=smoothness) return points_to_poses(new_points) def batch_visualize_tensor_global_pca(tensor_batch, num_components=3): B, C, H, W = tensor_batch.shape tensor_flat_all = tensor_batch.reshape(B, C, -1).permute(1, 0, 2).reshape(C, -1).T # StandardScaler + PCA computed on the tensor's device (the CPU version took ~8 s per video). # Same math as scikit-learn's covariance_eigh solver, which it uses for tall data like this, # including its sign convention (largest-magnitude entry of each component is positive). x = tensor_flat_all.double() n = x.shape[0] scale = x.std(dim=0, unbiased=False) scale[scale < 10 * torch.finfo(x.dtype).eps] = 1.0 x.sub_(x.mean(dim=0)).div_(scale) mean = x.mean(dim=0) cov = (x.T @ x - n * torch.outer(mean, mean)) / (n - 1) _, eigenvecs = torch.linalg.eigh(cov) components = eigenvecs.flip(1).T[:num_components] signs = torch.sign(components.gather(1, components.abs().argmax(dim=1, keepdim=True))) tensor_reduced_all = (x.sub_(mean) @ (components * signs).T).to(tensor_batch.dtype) del x tensor_reduced_all = tensor_reduced_all.T.reshape(num_components, B, H * W).permute(1, 0, 2) tensor_reduced_all = tensor_reduced_all - tensor_reduced_all.amin(dim=(1, 2), keepdim=True) tensor_reduced_all = tensor_reduced_all / tensor_reduced_all.amax(dim=(1, 2), keepdim=True) return tensor_reduced_all[:, :3].reshape(B, 3, H, W).float().cpu() def depth_to_colormap(depth_tensor, colormap='jet'): B, _, _, _ = depth_tensor.shape depth_tensor = (depth_tensor - depth_tensor.min()) / (depth_tensor.max() - depth_tensor.min()) depth_np = depth_tensor.squeeze(1).cpu().numpy() cmap = plt.get_cmap(colormap) colored_images = [] for i in range(B): colored_image = cmap(depth_np[i]) colored_images.append(colored_image[..., :3]) colored_tensor = torch.tensor(np.array(colored_images), dtype=torch.float32).permute(0, 3, 1, 2) return colored_tensor def save_video(frames, video_path, fps=24): clips = [mpy.ImageClip(frame).set_duration(1/fps) for frame in frames] video = mpy.concatenate_videoclips(clips, method="compose") video.write_videofile(video_path, fps=fps) def tensors_to_videos(all_images, all_depth_vis, all_fmap_vis, all_sems_vis, video_dir='videos', fps=24): B, C, H, W = all_images.shape assert all_depth_vis.shape == (B, C, H, W) assert all_fmap_vis.shape == (B, C, H, W) assert all_sems_vis is None or all_sems_vis.shape == (B, C, H, W) os.makedirs(video_dir, exist_ok=True) all_images = (all_images.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8) all_depth_vis = (all_depth_vis.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8) all_fmap_vis = (all_fmap_vis.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8) if all_sems_vis is not None: all_sems_vis = (all_sems_vis.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8) save_video(all_images, os.path.join(video_dir, 'output_images_video.mp4'), fps=fps) save_video(all_depth_vis, os.path.join(video_dir, 'output_depth_video.mp4'), fps=fps) save_video(all_fmap_vis, os.path.join(video_dir, 'output_fmap_video.mp4'), fps=fps) # save_video(all_sems_vis, os.path.join(video_dir, 'output_sems_video.mp4'), fps=fps) print(f'Videos saved to {video_dir}') def transfer_images_to_device(images, device): """ Transfer the loaded images to the specified device. Args: images (list): List of dictionaries containing image data. device (str or torch.device): The device to transfer the data to. Returns: list: List of dictionaries with image data transferred to the specified device. """ transferred_images = [] for img_dict in images: transferred_dict = { 'img': img_dict['img'].to(device), 'true_shape': torch.tensor(img_dict['true_shape'], device=device), 'idx': img_dict['idx'], 'instance': img_dict['instance'] } transferred_images.append(transferred_dict) return transferred_images def render_camera_path(video_poses, camera_params, gaussians, model, device, pipeline, bg_color, image_shape, with_semantics=True): """渲染相机路径的帮助函数 Args: video_poses: 相机位姿列表 camera_params: 包含extrinsics和intrinsics的相机参数 gaussians: 高斯模型 model: 特征提取模型 device: 计算设备 pipeline: 渲染管线 bg_color: 背景颜色 image_shape: 图像尺寸 Returns: rendered_images: 渲染的图像 rendered_feats: 渲染的特征图 rendered_depths: 渲染的深度图 rendered_sems: 渲染的语义图 """ extrinsics, intrinsics = camera_params rendered_images = [] rendered_feats = [] rendered_depths = [] rendered_sems = [] for i in range(len(video_poses)): target_extrinsics = torch.zeros(4, 4).to(device) target_extrinsics[3, 3] = 1.0 target_extrinsics[:3, :4] = torch.tensor(video_poses[i], device=device) camera = get_scaled_camera(extrinsics[0], target_extrinsics, intrinsics[0], 1.0, image_shape) rendered_output = render(camera, gaussians, pipeline, bg_color) rendered_images.append(rendered_output['render']) # 处理特征图 feature_map = rendered_output['feature_map'] feature_map = model.feature_expansion(feature_map[None, ...]) # 处理语义图 if with_semantics: logits = model.lseg_feature_extractor.decode_feature(feature_map, labelset=LABELS) semantic_map = torch.argmax(logits, dim=1) + 1 mask = COLORS[semantic_map.cpu()] mask = rearrange(mask, 'b h w c -> b c h w') rendered_sems.append(mask.squeeze(0)) # 降采样并上采样特征图 feature_map = feature_map[:, ::16, ...] feature_map = torch.nn.functional.interpolate(feature_map, scale_factor=2, mode='bilinear', align_corners=True) rendered_feats.append(feature_map[0]) del feature_map rendered_depths.append(rendered_output['depth']) # 堆叠并处理结果 rendered_images = torch.clamp(torch.stack(rendered_images, dim=0), 0, 1) rendered_feats = torch.stack(rendered_feats, dim=0) rendered_depths = torch.stack(rendered_depths, dim=0) rendered_sems = torch.stack(rendered_sems, dim=0) if with_semantics else None return rendered_images, rendered_feats, rendered_depths, rendered_sems @torch.no_grad() def render_video_from_file(file_list, model, output_path, device='cuda', resolution=224, n_interp=90, fps=30, path_type='default', render_input_path=True): """render_input_path=False skips the videos along the input-view path (only the ones in 'moved' are made).""" # 1. load images images = load_images(file_list, resolution, save_dir=os.path.join(output_path, 'processed_images')) images = transfer_images_to_device(images, device) # Transfer images to the specified device image_shape = images[0]['true_shape'][0] # 2. get camera pose pairs = make_pairs(images, prefilter=None, symmetrize=True) output = inference(pairs, model.mast3r, device, batch_size=1) mode = GlobalAlignerMode.PairViewer scene = global_aligner(output, device=device, mode=mode) extrinsics = scene.get_im_poses() intrinsics = scene.get_intrinsics() video_poses = generate_interpolated_path(extrinsics[:, :3, :].cpu().numpy(), n_interp=n_interp) # extrinsics: (b, 3, 4) # 3. get gaussians pred1, pred2 = model(*images) pred = merge_and_split_predictions(pred1, pred2) gaussians = GaussianModel.from_predictions(pred[0], sh_degree=3) # 4. 渲染原始视角 pipeline = DummyPipeline() bg_color = torch.tensor([0.0, 0.0, 0.0]).to(device) camera_params = (extrinsics, intrinsics) if render_input_path: rendered_images, rendered_feats, rendered_depths, rendered_sems = render_camera_path( video_poses, camera_params, gaussians, model, device, pipeline, bg_color, image_shape, with_semantics=False) # 5. 可视化 all_fmap_vis = batch_visualize_tensor_global_pca(rendered_feats) all_depth_vis = depth_to_colormap(rendered_depths) all_sems_vis = rendered_sems # 6. 保存视频 tensors_to_videos(rendered_images, all_depth_vis, all_fmap_vis, all_sems_vis, output_path, fps=fps) # 保存高斯点云 gaussians.save_ply(os.path.join(output_path, 'gaussians.ply')) # 7. 渲染移动后的视角 moved_extrinsics = move_c2w_along_z(extrinsics, 2.0) moved_video_poses = generate_interpolated_path(moved_extrinsics[:, :3, :].cpu().numpy(), n_interp=n_interp) camera_params = (extrinsics, intrinsics) moved_rendered_images, moved_rendered_feats, moved_rendered_depths, moved_rendered_sems = render_camera_path( moved_video_poses, camera_params, gaussians, model, device, pipeline, bg_color, image_shape, with_semantics=False) # 8. 可视化和保存移动后的结果 moved_all_fmap_vis = batch_visualize_tensor_global_pca(moved_rendered_feats) moved_all_depth_vis = depth_to_colormap(moved_rendered_depths) moved_all_sems_vis = moved_rendered_sems moved_output_path = os.path.join(output_path, 'moved') os.makedirs(moved_output_path, exist_ok=True) tensors_to_videos(moved_rendered_images, moved_all_depth_vis, moved_all_fmap_vis, moved_all_sems_vis, moved_output_path, fps=fps)