import os, re import json import numpy as np import cv2 import torch import imageio import pyexr import trimesh from PIL import Image from create_input import render_from_cameras_videos class DepthAlignMetric: """ 深度缩放与相机参数更新处理器 Attributes: moge_depth_dir (str): MOGe待处理深度目录 vggt_depth_dir (str): VGGT待处理深度目录 vggt_camera_json_file (str): VGGT关联的JSON文件目录 output_root (str): 输出根目录 """ def __init__(self, input_rgb_dir: str, moge_depth_dir: str, vggt_depth_dir: str, metric3d_depth_dir: str, vggt_camera_json_file: str, output_root: str): """ Args: moge_depth_dir: MOGe原始深度路径 vggt_depth_dir: VGGT原始深度路径 vggt_camera_json_file: VGGT关联JSON路径 output_root: 输出根目录，默认为./processed """ self.device = "cuda" if torch.cuda.is_available() else "cpu" # align depth and camera pose to metric level self.moge_depth_dir = moge_depth_dir self.vggt_depth_dir = vggt_depth_dir self.metric3d_depth_dir = metric3d_depth_dir self.vggt_camera_json_file = vggt_camera_json_file self.output_root = output_root # depth to pointmap self.metric_intrinsic = None self.metric_w2c = None self.input_rgb_dir = input_rgb_dir self.input_color_paths = [] # output depth / camera pose / pointmap self.output_metric_depth_dir = os.path.join(output_root, "output_metric_depth_dir") self.output_metric_camera_json = os.path.join(output_root, "output_metric_camera_json") self.output_metric_pointmap_dir = os.path.join(output_root, "output_metric_pointmap_dir") os.makedirs(self.output_metric_depth_dir, exist_ok=True) os.makedirs(self.output_metric_camera_json, exist_ok=True) os.makedirs(self.output_metric_pointmap_dir, exist_ok=True) def align_depth_scale(self): # align Moge depth to VGGT moge_align_depth_list, valid_mask_list = self.scale_moge_depth() # align moge depth and camera pose to metric depth self.align_metric_depth(moge_align_depth_list, valid_mask_list) def segment_sky_with_oneformer(self, image_path, skyseg_processor, skyseg_model, SKY_CLASS_ID, save_path=None): from PIL import Image image = Image.open(image_path) inputs = skyseg_processor(images=image, task_inputs=["semantic"], return_tensors="pt").to(skyseg_model.device) with torch.no_grad(): outputs = skyseg_model(**inputs) # 获取语义分割结果 predicted_semantic_map = skyseg_processor.post_process_semantic_segmentation(outputs, \ target_sizes=[image.size[::-1]])[0] # 提取天空区域 sky_mask = (predicted_semantic_map == SKY_CLASS_ID).cpu().numpy().astype(np.uint8) * 255 # erosion sky kernel = np.ones((3,3), np.uint8) sky_mask = cv2.erode(sky_mask, kernel, iterations=1) # 如果需要保存 if save_path: cv2.imwrite(save_path, sky_mask) return sky_mask def get_valid_depth(self, vggt_files, moge_files, input_rgb_files, skyseg_processor, skyseg_model, SKY_CLASS_ID): moge_align_depth_list = [] valid_mask_list = [] all_valid_max_list = [] for vggt_file, moge_file, input_rgb_file in zip(vggt_files, moge_files, input_rgb_files): # 读取深度数据 depth_moge = pyexr.read(os.path.join(self.moge_depth_dir, moge_file)).squeeze() depth_vggt = pyexr.read(os.path.join(self.vggt_depth_dir, vggt_file)).squeeze() depth_vggt = cv2.resize(depth_vggt, dsize=(depth_moge.shape[1], depth_moge.shape[0]), \ interpolation=cv2.INTER_LINEAR) depth_vggt = torch.from_numpy(depth_vggt).float().to(self.device) depth_moge = torch.from_numpy(depth_moge).float().to(self.device) # segmentation sky sky_ima_path = os.path.join(self.input_rgb_dir, input_rgb_file) sky_mask = self.segment_sky_with_oneformer(sky_ima_path, skyseg_processor, skyseg_model, SKY_CLASS_ID) sky_mask_tensor = torch.from_numpy(sky_mask).float().to(self.device) sky_mask = (sky_mask_tensor > 0) # 天空区域为True valid_masks = ( # (H, W) torch.isfinite(depth_moge) & (depth_moge > 0) & torch.isfinite(depth_vggt) & (depth_vggt > 0) & ~sky_mask # 非天空区域 ) # depth_moge 无效部分 设置为 有效部分最大值的1.5倍 避免final_align_depth出现负数 depth_moge[~valid_masks] = depth_moge[valid_masks].max() * 1 source_inv_depth = 1.0 / depth_moge target_inv_depth = 1.0 / depth_vggt # print(f'倒数值:{source_inv_depth.min()}, {source_inv_depth.max()}') # 0.03 ～ 2.2 source_mask, target_mask = valid_masks, valid_masks # Remove outliers 2/8分最合适 outlier_quantiles = torch.tensor([0.2, 0.8], device=self.device) source_data_low, source_data_high = torch.quantile( source_inv_depth[source_mask], outlier_quantiles ) target_data_low, target_data_high = torch.quantile( target_inv_depth[target_mask], outlier_quantiles ) source_mask = (source_inv_depth > source_data_low) & ( source_inv_depth < source_data_high ) target_mask = (target_inv_depth > target_data_low) & ( target_inv_depth < target_data_high ) mask = torch.logical_and(source_mask, target_mask) mask = torch.logical_and(mask, valid_masks) source_data = source_inv_depth[mask].view(-1, 1) target_data = target_inv_depth[mask].view(-1, 1) ones = torch.ones((source_data.shape[0], 1), device=self.device) source_data_h = torch.cat([source_data, ones], dim=1) transform_matrix = torch.linalg.lstsq(source_data_h, target_data).solution scale, bias = transform_matrix[0, 0], transform_matrix[1, 0] aligned_inv_depth = source_inv_depth * scale + bias valid_inv_depth = aligned_inv_depth > 0 # 创建新的有效掩码 valid_masks = valid_masks & valid_inv_depth # 合并到原有效掩码 valid_mask_list.append(valid_masks) final_align_depth = 1.0 / aligned_inv_depth moge_align_depth_list.append(final_align_depth) all_valid_max_list.append(final_align_depth[valid_masks].max().item()) return moge_align_depth_list, valid_mask_list, all_valid_max_list def scale_moge_depth(self): vggt_files = sorted(f for f in os.listdir(self.vggt_depth_dir) if f.endswith('.exr')) moge_files = sorted(f for f in os.listdir(self.moge_depth_dir) if f.endswith('.exr')) input_rgb_files = sorted(f for f in os.listdir(self.input_rgb_dir) if f.endswith('.png')) if len(vggt_files) != len(moge_files): raise ValueError("文件数量不匹配") from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation skyseg_processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large") skyseg_model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/oneformer_coco_swin_large") skyseg_model.to(self.device) # 定义天空类别的ID 119 SKY_CLASS_ID = 119 moge_align_depth_list, valid_mask_list, all_valid_max_list = self.get_valid_depth( vggt_files, moge_files, input_rgb_files, skyseg_processor, skyseg_model, SKY_CLASS_ID ) # 计算所有帧的有效最大值的中位数 valid_max_array = np.array(all_valid_max_list) q50 = np.quantile(valid_max_array, 0.50) # 计算50%分位点 filtered_max = valid_max_array[valid_max_array <= q50] # 过滤超过分位点的异常值 # 取过滤后数据的最大值（正常范围内的最大值） global_avg_max = np.max(filtered_max) max_sky_value = global_avg_max * 5 max_sky_value = np.minimum(max_sky_value, 1000) # 相对深度最远不能超过 1000 # 统一设置所有帧的无效区域值 for i, (moge_depth, valid_mask) in enumerate(zip(moge_align_depth_list, valid_mask_list)): moge_depth[~valid_mask] = max_sky_value # 统计超限点占比（在clamp之前） over_count = torch.sum(moge_depth > max_sky_value).item() total_pixels = moge_depth.numel() over_ratio = over_count / total_pixels * 100 moge_depth = torch.clamp(moge_depth, max=max_sky_value) moge_align_depth_list[i] = moge_depth # 更新处理后的深度图 return moge_align_depth_list, valid_mask_list def align_metric_depth(self, moge_align_depth_list, valid_mask_list): # 获取metric文件列表 metric_files = sorted(f for f in os.listdir(self.metric3d_depth_dir) if f.endswith('.exr')) metric_scales_list = [] # 遍历所有深度图对 for idx, (metric_file, moge_depth) in enumerate(zip(metric_files, moge_align_depth_list)): depth_metric3d = pyexr.read(os.path.join(self.metric3d_depth_dir, metric_file)).squeeze() depth_metric3d = torch.from_numpy(depth_metric3d).float().to(self.device) # 获取对应帧的掩码 valid_mask = valid_mask_list[idx].to(self.device) # 提取有效区域数据 valid_metric = depth_metric3d[valid_mask] valid_moge = moge_depth[valid_mask] # 分位数差计算 metric_diff = torch.quantile(valid_metric, 0.8) - torch.quantile(valid_metric, 0.2) moge_diff = torch.quantile(valid_moge, 0.8) - torch.quantile(valid_moge, 0.2) metric_scale = metric_diff / moge_diff metric_scales_list.append(metric_scale.cpu().numpy()) # 计算全局平均缩放因子 metric_scales_mean = np.mean(metric_scales_list) # 应用全局缩放 保存 metric depth for idx, (metric_file, moge_depth) in enumerate(zip(metric_files, moge_align_depth_list)): metric_moge_depth = (moge_depth * metric_scales_mean).cpu().numpy() # 保存深度文件 output_path = os.path.join( self.output_metric_depth_dir, f"{os.path.splitext(metric_file)[0]}_metric.exr" ) pyexr.write(output_path, metric_moge_depth, channel_names=["Y"]) # 阶段3：更新相机参数 with open(self.vggt_camera_json_file, 'r') as f: camera_data = json.load(f) # 更新所有帧的平移分量 for frame_info in camera_data.values(): w2c_matrix = np.array(frame_info['w2c']) w2c_matrix[:3, 3] *= metric_scales_mean # 直接使用计算好的全局缩放因子 frame_info['w2c'] = w2c_matrix.tolist() # 保存更新后的相机参数 output_json_path = os.path.join( self.output_metric_camera_json, os.path.basename(self.vggt_camera_json_file) ) with open(output_json_path, 'w') as f: json.dump(camera_data, f, indent=4) def load_metirc_camera_parameters(self): # 修改：增加color_dir参数 metric_camera_json = os.path.join(self.output_metric_camera_json, 'colmap_data.json') with open(metric_camera_json, 'r') as f: data = json.load(f) # load metric camera parameters sorted_frames = sorted(data.items(), key=lambda x: int(x[0])) first_frame_key, first_frame_data = sorted_frames[0] self.metric_intrinsic = [np.array(frame['intrinsic']) for frame in data.values()] self.metric_w2c = [np.array(frame['w2c']) for frame in data.values()] # 加载pointmap input rgb 文件路径 self.input_color_paths = sorted( [os.path.join(self.input_rgb_dir, f) for f in os.listdir(self.input_rgb_dir) if f.endswith(".png")], key=lambda x: int(os.path.basename(x).split("_")[1].split(".")[0]) ) def depth_to_pointmap(self): num_frames = len(self.metric_w2c) for frame_index in range(num_frames): exr_path = os.path.join(self.output_metric_depth_dir, f"frame_{frame_index+1:05d}_metric.exr") depth_data = pyexr.read(exr_path).squeeze() depth_tensor = torch.from_numpy(depth_data).to(self.device, torch.float32) # 生成点云 height, width = depth_tensor.shape K_tensor = torch.from_numpy(self.metric_intrinsic[frame_index]).to(device=self.device, dtype=torch.float32) w2c = torch.from_numpy(self.metric_w2c[frame_index]).to(device=self.device, dtype=torch.float32) camtoworld = torch.inverse(w2c) # 生成相机坐标系坐标 u = torch.arange(width, device=self.device).float() v = torch.arange(height, device=self.device).float() u_grid, v_grid = torch.meshgrid(u, v, indexing='xy') fx, fy = K_tensor[0, 0], K_tensor[1, 1] cx, cy = K_tensor[0, 2], K_tensor[1, 2] x_cam = (u_grid - cx) * depth_tensor / fx y_cam = (v_grid - cy) * depth_tensor / fy z_cam = depth_tensor cam_coords_points = torch.stack([x_cam, y_cam, z_cam], dim=-1) R_cam_to_world = camtoworld[:3, :3] t_cam_to_world = camtoworld[:3, 3] world_coords_points = torch.matmul(cam_coords_points, R_cam_to_world.T) + t_cam_to_world # # 保存带颜色的点云 color_numpy = np.array(Image.open(self.input_color_paths[frame_index])) # 读取为HWC colors_rgb = color_numpy.reshape(-1, 3) # 转换回HWC并展平 vertices_3d = world_coords_points.reshape(-1, 3).cpu().numpy() point_cloud_data = trimesh.PointCloud(vertices=vertices_3d, colors=colors_rgb) point_cloud_data.export(f"{self.output_metric_pointmap_dir}/pcd_{frame_index+1:04d}.ply") # 保存为pointmap npy pointmap_data = world_coords_points.cpu().numpy() np.save(f"{self.output_metric_pointmap_dir}/pointmap_{frame_index+1:04d}.npy", pointmap_data) def render_from_cameras(self): render_output_dir = os.path.join(self.output_root, "rendered_views") os.makedirs(render_output_dir, exist_ok=True) select_frame = 0 npy_files = sorted( [f for f in os.listdir(self.output_metric_pointmap_dir) if f.endswith(".npy")], key=lambda x: int(re.findall(r'\d+', x)[0]) ) npy_path = os.path.join(self.output_metric_pointmap_dir, npy_files[select_frame]) # 读取npy_path pointmap = np.load(npy_path) points = pointmap.reshape(-1, 3) color_numpy = np.array(Image.open(self.input_color_paths[select_frame])) # 读取为HWC colors_rgb = color_numpy.reshape(-1, 3) # 转换回HWC并展平 colors = colors_rgb[:, :3] height, width = cv2.imread(self.input_color_paths[0]).shape[:2] renders, masks, _ = render_from_cameras_videos( points, colors, self.metric_w2c, self.metric_intrinsic, height, width ) # 使用imageio保存所有结果 for i, (render, mask) in enumerate(zip(renders, masks)): # 保存渲染图 render_path = os.path.join(render_output_dir, f"render_{i:04d}.png") imageio.imwrite(render_path, render) # 保存掩码图 mask_path = os.path.join(render_output_dir, f"mask_{i:04d}.png") imageio.imwrite(mask_path, mask) print(f"All results saved to: {render_output_dir}") if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Depth alignment and metric processing.") parser.add_argument('--image_dir', type=str, required=True, help='Input RGB directory') parser.add_argument('--moge_depth_dir', type=str, required=True, help='MOGe depth directory') parser.add_argument('--vggt_depth_dir', type=str, required=True, help='VGGT depth directory') parser.add_argument('--metric3d_depth_dir', type=str, required=True, help='Metric3D depth directory') parser.add_argument('--vggt_camera_json_file', type=str, required=True, help='VGGT camera JSON file') parser.add_argument('--output_dir', type=str, required=True, help='Output root directory') args = parser.parse_args() depth_align_processor = DepthAlignMetric( input_rgb_dir=args.image_dir, moge_depth_dir=args.moge_depth_dir, vggt_depth_dir=args.vggt_depth_dir, metric3d_depth_dir=args.metric3d_depth_dir, vggt_camera_json_file=args.vggt_camera_json_file, output_root=args.output_dir ) depth_align_processor.align_depth_scale() depth_align_processor.load_metirc_camera_parameters() depth_align_processor.depth_to_pointmap() depth_align_processor.render_from_cameras()