File size: 6,187 Bytes
ed66dc1 | 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 | """CONTRACT.md line (b) renderer: a 3DGS .ply -> 4 fixed views, 512², on the Mac CPU.
Projection, covariance and tile intersection are gsplat 1.5.3's pure-PyTorch functions (nerfstudio-project/gsplat,
Apache-2.0, `gsplat/cuda/_torch_impl.py`, loaded from the installed wheel, unmodified). Compositing follows gsplat's
`accumulate`: alpha = min(0.999, opacity * exp(-sigma)), weight = alpha * prod(1 - alpha in front), over every
Gaussian whose 3.33-sigma box touches the 16x16 tile, front to back by camera depth.
python render_views.py --ply A.ply [--ply B.ply] --out DIR [--psnr] (with two .ply: PSNR of B vs A + side-by-side)
"""
import argparse
import glob
import importlib.util
import json
import math
import os
import sys
import numpy as np
import torch
from PIL import Image
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import idol_glue as GL # noqa: E402
_p = glob.glob(os.path.join(os.path.dirname(torch.__file__), '..', 'gsplat', 'cuda', '_torch_impl.py'))[0]
_spec = importlib.util.spec_from_file_location('gsplat_torch_impl', _p)
GS = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(GS)
C0 = 0.28209479177387814
SIZE, TILE = 512, 16
AZIMUTHS = (0.0, 90.0, 180.0, 270.0)
ELEVATION, DISTANCE, FOV_Y = 20.0, 2.2, 40.0
def cameras():
viewmats, Ks = [], []
# idol's PLY is y-down, as 3DGS viewers read it (the object's top is at -y): image up = -y, cameras above at -y
up = np.array([0.0, -1.0, 0.0])
f = (SIZE / 2) / math.tan(math.radians(FOV_Y) / 2)
for az in AZIMUTHS:
a, e = math.radians(az), math.radians(ELEVATION)
c = DISTANCE * np.array([math.cos(e) * math.sin(a), -math.sin(e), math.cos(e) * math.cos(a)])
z = -c / np.linalg.norm(c) # forward (OpenCV camera: x right, y down, z forward)
x = np.cross(z, up); x /= np.linalg.norm(x)
y = np.cross(z, x)
R = np.stack([x, y, z])
m = np.eye(4); m[:3, :3] = R; m[:3, 3] = -R @ c
viewmats.append(m)
Ks.append(np.array([[f, 0, SIZE / 2], [0, f, SIZE / 2], [0, 0, 1.0]]))
return torch.tensor(np.stack(viewmats), dtype=torch.float32), torch.tensor(np.stack(Ks), dtype=torch.float32)
def load(path):
names, rows = GL.read_ply(path)
col = {n: i for i, n in enumerate(names)}
t = torch.from_numpy(np.ascontiguousarray(rows))
means = t[:, [col['x'], col['y'], col['z']]]
quats = torch.nn.functional.normalize(t[:, [col['rot_0'], col['rot_1'], col['rot_2'], col['rot_3']]], dim=-1)
scales = torch.exp(t[:, [col['scale_0'], col['scale_1'], col['scale_2']]])
opac = torch.sigmoid(t[:, col['opacity']])
colors = (0.5 + C0 * t[:, [col['f_dc_0'], col['f_dc_1'], col['f_dc_2']]]).clamp_min(0.0)
return means, quats, scales, opac, colors
@torch.no_grad()
def render(path):
means, quats, scales, opac, colors = load(path)
viewmats, Ks = cameras()
covars, _ = GS._quat_scale_to_covar_preci(quats, scales, compute_covar=True, compute_preci=False)
radii, means2d, depths, conics, _ = GS._fully_fused_projection(means, covars, viewmats, Ks, SIZE, SIZE,
eps2d=0.3, near_plane=0.01, far_plane=1e10)
tw = th = SIZE // TILE
_, isect_ids, flatten_ids = GS._isect_tiles(means2d, radii, depths, TILE, tw, th, sort=True)
offsets = GS._isect_offset_encode(isect_ids, len(AZIMUTHS), tw, th).flatten().tolist() + [len(flatten_ids)]
N = means.shape[0]
flat = flatten_ids.long()
out = torch.zeros(len(AZIMUTHS), SIZE, SIZE, 3)
py, px = torch.meshgrid(torch.arange(TILE, dtype=torch.float32), torch.arange(TILE, dtype=torch.float32),
indexing='ij')
for k in range(len(offsets) - 1):
s, e = offsets[k], offsets[k + 1]
if e <= s:
continue
cam, rem = divmod(k, tw * th)
ty, tx = divmod(rem, tw)
gid = flat[s:e] - cam * N
m2 = means2d[cam, gid]
cn = conics[cam, gid]
dx = (px.flatten() + tx * TILE + 0.5)[None, :] - m2[:, 0:1]
dy = (py.flatten() + ty * TILE + 0.5)[None, :] - m2[:, 1:2]
sigma = 0.5 * (cn[:, 0:1] * dx * dx + cn[:, 2:3] * dy * dy) + cn[:, 1:2] * dx * dy
alpha = torch.clamp_max(opac[gid][:, None] * torch.exp(-sigma), 0.999)
trans = torch.cumprod(torch.cat([torch.ones(1, alpha.shape[1]), 1 - alpha[:-1]], dim=0), dim=0)
w = alpha * trans
rgb = (w.t() @ colors[gid]).reshape(TILE, TILE, 3)
out[cam, ty * TILE:(ty + 1) * TILE, tx * TILE:(tx + 1) * TILE] = rgb
return out.clamp(0, 1)
def psnr(a, b):
mse = float(((a.double() - b.double()) ** 2).mean())
return 10 * math.log10(1.0 / max(mse, 1e-12)), mse
def to_u8(img):
return (img.numpy() * 255.0 + 0.5).clip(0, 255).astype(np.uint8)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--ply', action='append', required=True, help='first = reference (author), second = ours')
ap.add_argument('--out', required=True)
ap.add_argument('--tag', default='')
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)
views = [render(p) for p in a.ply]
for i, v in enumerate(views):
for j, az in enumerate(AZIMUTHS):
Image.fromarray(to_u8(v[j])).save(os.path.join(a.out, f'view{i}_az{int(az):03d}.png'))
res = {'ply': a.ply, 'cameras': {'azimuth_deg': AZIMUTHS, 'elevation_deg': ELEVATION, 'distance': DISTANCE,
'fov_y_deg': FOV_Y, 'size': SIZE}, 'renderer': 'gsplat 1.5.3 _torch_impl'}
if len(views) == 2:
p_all, mse = psnr(views[0], views[1])
res.update(psnr_db=p_all, mse=mse, psnr_per_view_db=[psnr(views[0][j], views[1][j])[0]
for j in range(len(AZIMUTHS))])
grid = np.concatenate([np.concatenate([to_u8(v[j]) for j in range(len(AZIMUTHS))], axis=1) for v in views])
Image.fromarray(grid).save(os.path.join(a.out, f'b_{a.tag or "pair"}.png'))
json.dump(res, open(os.path.join(a.out, 'render.json'), 'w'), indent=1)
print(json.dumps(res), flush=True)
if __name__ == '__main__':
main()
|