File size: 6,543 Bytes
23e73f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | 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
|