| 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) |
| 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 |
| 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 |
|
|
