import os import sys import math import argparse import torch import numpy as np from PIL import Image, ImageOps from plyfile import PlyData import gsplat sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') @torch.no_grad() def load_ply(path): plydata = PlyData.read(path) v = plydata['vertex'] means = np.stack((v['x'], v['y'], v['z']), axis=-1) quats = np.stack((v['rot_0'], v['rot_1'], v['rot_2'], v['rot_3']), axis=-1) scales = np.stack((v['scale_0'], v['scale_1'], v['scale_2']), axis=-1) opacities = v['opacity'] f_dc = np.stack((v['f_dc_0'], v['f_dc_1'], v['f_dc_2']), axis=-1) f_rest = np.stack([v[f'f_rest_{i}'] for i in range(45)], axis=-1) device = torch.device("cuda") means = torch.tensor(means, dtype=torch.float32, device=device) quats = torch.tensor(quats, dtype=torch.float32, device=device) scales = torch.exp(torch.tensor(scales, dtype=torch.float32, device=device)) opacities = torch.sigmoid(torch.tensor(opacities, dtype=torch.float32, device=device)) f_dc = torch.tensor(f_dc, dtype=torch.float32, device=device).unsqueeze(1) f_rest = torch.tensor(f_rest, dtype=torch.float32, device=device) f_rest = f_rest.view(-1, 3, 15).transpose(1, 2) shs = torch.cat([f_dc, f_rest], dim=1) return means, quats, scales, opacities, shs def load_test_cameras(source_path, resolution): from scene.dataset_readers import readColmapSceneInfo, readNerfSyntheticInfo from utils.graphics_utils import getWorld2View2 import inspect parsed = [] is_synthetic = os.path.exists(os.path.join(source_path, "transforms_test.json")) and not os.path.exists(os.path.join(source_path, "sparse")) if not is_synthetic: img_dir = f"images_{resolution}" if resolution > 1 and os.path.exists(os.path.join(source_path, f"images_{resolution}")) else "images" sig = inspect.signature(readColmapSceneInfo) args_list = [] for i, (k, p) in enumerate(sig.parameters.items()): if i == 0: args_list.append(source_path) elif i == 1: args_list.append(img_dir) elif k == "eval": args_list.append(True) # <--- 核心修改：让 Native 决定谁是 Test 集 elif k == "train_test_exp": args_list.append(False) else: args_list.append(p.default if p.default != inspect.Parameter.empty else "") scene_info = readColmapSceneInfo(*args_list) test_cams = scene_info.test_cameras # 直接沿用无损的 Native Test 阵列 else: sig = inspect.signature(readNerfSyntheticInfo) args_list = [] for i, (k, p) in enumerate(sig.parameters.items()): if i == 0: args_list.append(source_path) elif k == "eval": args_list.append(True) elif k == "extension": args_list.append(".png") else: args_list.append(p.default if p.default != inspect.Parameter.empty else "") scene_info = readNerfSyntheticInfo(*args_list) test_cams = scene_info.test_cameras for c in test_cams: w = int(round(c.width / resolution)) h = int(round(c.height / resolution)) viewmat = getWorld2View2(np.array(c.R), np.array(c.T)).astype(np.float32) fx = w / (2 * math.tan(c.FovX / 2)) fy = h / (2 * math.tan(c.FovY / 2)) K = np.array([[fx, 0, w/2], [0, fy, h/2], [0, 0, 1]], dtype=np.float32) pil_img = getattr(c, 'image', None) if pil_img is None: img_path = getattr(c, 'image_path', None) if not img_path or not os.path.exists(img_path): folder = f"images_{resolution}" if resolution > 1 and os.path.exists(os.path.join(source_path, f"images_{resolution}")) else "images" img_path = os.path.join(source_path, folder, c.image_name) pil_img = Image.open(img_path) pil_img = ImageOps.exif_transpose(pil_img) parsed.append({ 'name': c.image_name, 'pil_image': pil_img, 'viewmat': viewmat, 'K': K, 'width': w, 'height': h, }) return parsed def compute_psnr(a, b): mse = torch.mean((a - b) ** 2) return 100.0 if mse == 0 else (10 * torch.log10(1.0 / mse)).item() def main(): parser = argparse.ArgumentParser() parser.add_argument("--ply_path", required=True) parser.add_argument("--source_path", required=True) parser.add_argument("--model_path", required=True) parser.add_argument("--output_dir", required=True) parser.add_argument("--resolution", type=int, default=2) parser.add_argument("--bg_color", type=str, default="0,0,0") args = parser.parse_args() from utils.general_utils import PILtoTorch device = torch.device("cuda") means, quats, scales, opacities, shs = load_ply(args.ply_path) cameras = load_test_cameras(args.source_path, args.resolution) bg_parts = [float(x) for x in args.bg_color.split(",")] bg_color = torch.tensor(bg_parts, dtype=torch.float32, device=device) bg_color_chw = bg_color.view(3, 1, 1) psnrs = [] for i, cam in enumerate(cameras): viewmat = torch.tensor(cam['viewmat'], device=device).unsqueeze(0) K = torch.tensor(cam['K'], device=device).unsqueeze(0) colors, _, _ = gsplat.rasterization( means=means, quats=quats, scales=scales, opacities=opacities, colors=shs, viewmats=viewmat, Ks=K, width=cam['width'], height=cam['height'], sh_degree=3, packed=True, render_mode='RGB', backgrounds=bg_color.unsqueeze(0), near_plane=0.01, far_plane=100.0, ) render_img = colors[0].clamp(0, 1) gt_chw_full = PILtoTorch(cam['pil_image'], (cam['width'], cam['height'])).to(device) if gt_chw_full.shape[0] == 4: alpha = gt_chw_full[3:4, ...] rgb = gt_chw_full[:3, ...] gt_chw = rgb * alpha + bg_color_chw * (1.0 - alpha) else: gt_chw = gt_chw_full[:3, ...] gt_img = gt_chw.permute(1, 2, 0) psnr_val = compute_psnr(render_img, gt_img) psnrs.append(psnr_val) if not psnrs: print("\nRENDER_SINGLE DONE — STATUS: FAIL — DELTA: nan dB") return mean_psnr = float(np.mean(psnrs)) baseline_path = os.path.join(args.model_path, "metrics_test_iter30000.json") baseline = 32.4354