| """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 |
|
|
| _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 = [], [] |
| |
| 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) |
| 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() |
|
|
