# -*- coding: utf-8 -*- """ Minimal NeRF (PyTorch, CUDA) — course-project implementation. Reads a nerfstudio-style transforms.json (OpenCV/COLMAP convention, c2w) and trains a vanilla NeRF (positional encoding + 8x256 MLP + hierarchical sampling). Then renders novel/test views and exports a dense coloured point cloud by unprojecting rendered depth over all training views. Usage: python nerf_train.py --transforms output/colmap/cameras_export/transforms.json \ --img_dir mvs_project/images --outdir output/nerf --steps 30000 """ import os, sys, time, json, argparse import numpy as np import torch import torch.nn as nn import torch.nn.functional as F # ---------------------------------------------------------------- positional encoding def posenc(x, L): out = [x] for i in range(L): for fn in (torch.sin, torch.cos): out.append(fn((2.0 ** i) * torch.pi * x)) return torch.cat(out, dim=-1) # ---------------------------------------------------------------- NeRF MLP class NeRF(nn.Module): def __init__(self, D=8, W=256, in_xyz=63, in_dir=27, skip=4): super().__init__() self.D, self.W, self.skip = D, W, skip self.in_xyz, self.in_dir = in_xyz, in_dir for i in range(D): if i == 0: layer = nn.Linear(in_xyz, W) elif i == skip: layer = nn.Linear(W + in_xyz, W) else: layer = nn.Linear(W, W) setattr(self, f"xyz_{i+1}", layer) self.xyz_out = nn.Linear(W, 1 + W) # sigma + bottleneck self.dir_1 = nn.Linear(W + in_dir, W // 2) # view dependence self.rgb_out = nn.Linear(W // 2, 3) def forward(self, x, d): xyz_ = posenc(x, 10) dir_ = posenc(d, 4) h = xyz_ for i in range(self.D): if i == self.skip: h = torch.cat([h, xyz_], dim=-1) h = F.relu(getattr(self, f"xyz_{i+1}")(h)) out = self.xyz_out(h) # [N, 1+W] sigma = F.softplus(out[..., :1] - 1.0) # density feat = out[..., 1:] h2 = F.relu(self.dir_1(torch.cat([feat, dir_], dim=-1))) rgb = torch.sigmoid(self.rgb_out(h2)) return rgb, sigma.squeeze(-1) # ---------------------------------------------------------------- ray helpers def get_rays(c2w, fx, fy, cx, cy, H, W, uv): """uv: [N,2] pixel coords (u,v) float. Returns rays_o, rays_d in world (OpenCV).""" x = (uv[:, 0] - cx) / fx y = (uv[:, 1] - cy) / fy dirs_cam = torch.stack([x, y, torch.ones_like(x)], dim=-1) # z forward dirs_cam = dirs_cam / torch.norm(dirs_cam, dim=-1, keepdim=True) R = c2w[:3, :3].to(x.device) rays_d = dirs_cam @ R.T rays_o = c2w[:3, 3].to(x.device).expand(rays_d.shape[0], -1) return rays_o, rays_d def sample_along_rays(rays_o, rays_d, near, far, N, perturb=True): """stratified sampling -> points + z_vals""" N_rays = rays_o.shape[0] t = torch.linspace(0.0, 1.0, N, device=rays_o.device)[None, :] z_vals = near * (1 - t) + far * t # [1,N] if perturb: mids = 0.5 * (z_vals[:, 1:] + z_vals[:, :-1]) upper = torch.cat([mids, z_vals[:, -1:]], dim=-1) lower = torch.cat([z_vals[:, :1], mids], dim=-1) z_vals = lower + (upper - lower) * torch.rand(N_rays, N, device=rays_o.device) pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., None] # [R,N,3] return pts, z_vals def raw2outputs(raw, z_vals, rays_d, white_bkgd=False): """volume rendering. raw: [R, N, 4] = [rgb(3), sigma]""" sig = F.relu(raw[..., 3]) dists = z_vals[..., 1:] - z_vals[..., :-1] dists = torch.cat([dists, torch.full_like(dists[..., :1], 1e10)], dim=-1) alpha = 1.0 - torch.exp(-sig * dists) T = torch.cumprod(1.0 - alpha + 1e-10, dim=-1) T = torch.cat([torch.ones_like(T[..., :1]), T[..., :-1]], dim=-1) w = T * alpha # weights rgb = torch.sum(w[..., None] * raw[..., :3], dim=-2) depth = torch.sum(w * z_vals, dim=-1) acc = torch.sum(w, dim=-1) return rgb, depth, acc, w def render_rays(model, rays_o, rays_d, near, far, Nc, Nf, perturb=True): """hierarchical sampling, returns coarse + fine outputs.""" pts, z_vals = sample_along_rays(rays_o, rays_d, near, far, Nc, perturb) raw_c = run_network(model, pts, rays_d) # [R,N,4] rgb_c, depth_c, acc_c, w_c = raw2outputs(raw_c, z_vals, rays_d) # fine sampling from coarse weights w = w_c[:, :-1] + 1e-6 w = w / w.sum(dim=-1, keepdim=True) R = rays_o.shape[0] z_vals_R = z_vals.expand(R, -1) # [R, Nc] mids = 0.5 * (z_vals_R[:, 1:] + z_vals_R[:, :-1]) # [R, Nc-1] inds = torch.multinomial(w, Nf, replacement=True) inds = inds.sort(dim=-1).values z_fine = torch.gather(mids, -1, inds) # [R, Nf] z_all, _ = torch.sort(torch.cat([z_vals_R, z_fine], dim=-1), dim=-1) pts_fine = rays_o[:, None, :] + rays_d[:, None, :] * z_all[..., None] raw_f = run_network(model, pts_fine, rays_d) rgb_f, depth_f, acc_f, w_f = raw2outputs(raw_f, z_all, rays_d) return (rgb_c, depth_c, acc_c), (rgb_f, depth_f, acc_f) def run_network(model, pts, rays_d): flat = pts.reshape(-1, 3) d = rays_d[:, None, :].expand(pts.shape).reshape(-1, 3) rgb, sig = model(flat, d) N = pts.shape[1] rgb = rgb.reshape(-1, N, 3) sig = sig.reshape(-1, N) return torch.cat([rgb, sig[..., None]], dim=-1) # ---------------------------------------------------------------- image loading def load_image(path): from PIL import Image img = np.array(Image.open(path).convert("RGB")).astype(np.float32) / 255.0 return img # ---------------------------------------------------------------- data def load_dataset(transforms_path, img_dir): tr = json.load(open(transforms_path)) fx, fy = tr["fl_x"], tr["fl_y"] cx, cy = tr["cx"], tr["cy"] H, W = tr["h"], tr["w"] imgs, c2ws, names = [], [], [] for fr in tr["frames"]: base = os.path.basename(fr["file_path"]) p = os.path.join(img_dir, base) imgs.append(load_image(p)) c2ws.append(np.asarray(fr["transform_matrix"], np.float32)) names.append(base) imgs = np.stack(imgs).astype(np.float32) # [N,H,W,3] c2ws = np.stack(c2ws) # [N,4,4] return imgs, c2ws, names, (fx, fy, cx, cy, H, W) def sample_batch(imgs, c2ws, intr, batch, device): fx, fy, cx, cy, H, W = intr n = len(imgs) idx = np.random.randint(0, n, batch) u = np.random.rand(batch) * W v = np.random.rand(batch) * H ui = np.clip(u.astype(int), 0, W - 1) vi = np.clip(v.astype(int), 0, H - 1) cols = imgs[idx, vi, ui] # [batch,3] cols = torch.from_numpy(cols).to(device) c2w = c2ws[idx] # [batch,4,4] R = c2w[:, :3, :3] t = c2w[:, :3, 3] x = (u - cx) / fx y = (v - cy) / fy dirs_cam = np.stack([x, y, np.ones_like(x)], -1) dirs_cam /= np.linalg.norm(dirs_cam, axis=-1, keepdims=True) rays_d = np.einsum('bij,bj->bi', R, dirs_cam).astype(np.float32) rays_o = t return (torch.from_numpy(rays_o).to(device), torch.from_numpy(rays_d).to(device), cols) # ---------------------------------------------------------------- point cloud export @torch.no_grad() def export_pointcloud(model, imgs, c2ws, names, intr, out_path, near, far, Nc, Nf, device, step=2, thresh=0.5): from PIL import Image fx, fy, cx, cy, H, W = intr all_pts, all_col = [], [] model.eval() for vi, (c2w, name) in enumerate(zip(c2ws, names)): # pixel grid at stride `step` us = torch.arange(0, W, step, device=device).float() + 0.5 * step vs = torch.arange(0, H, step, device=device).float() + 0.5 * step uv = torch.stack(torch.meshgrid(us, vs, indexing="ij"), dim=-1).reshape(-1, 2) c2w_t = torch.from_numpy(c2w).to(device) P = [] for s in range(0, uv.shape[0], 4096): uvb = uv[s:s+4096] o, d = get_rays(c2w_t, fx, fy, cx, cy, H, W, uvb) (_, _, _), (rgb, depth, acc) = render_rays(model, o, d, near, far, Nc, Nf, perturb=False) pts = o + depth[:, None] * d keep = acc > thresh P.append((pts[keep], rgb[keep])) if not P: continue pts = torch.cat([p[0] for p in P], 0).cpu().numpy() col = torch.cat([p[1] for p in P], 0).cpu().numpy() all_pts.append(pts); all_col.append(col) if (vi + 1) % 12 == 0: print(f" extracted view {vi+1}/{len(names)}", flush=True) pts = np.concatenate(all_pts, 0) col = np.concatenate(all_col, 0) print(f"total points: {pts.shape[0]}") # voxel downsample pts, col = voxel_downsample(pts, col, voxel=0.01) print(f"after voxel downsample (1cm): {pts.shape[0]}") write_ply(out_path, pts, col) def voxel_downsample(pts, col, voxel=0.01): keys = np.floor(pts / voxel).astype(np.int64) _, idx = np.unique(keys, axis=0, return_index=True) return pts[idx], col[idx] def write_ply(path, pts, col): os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "w") as f: f.write("ply\nformat ascii 1.0\n") f.write(f"element vertex {len(pts)}\n") f.write("property float x\nproperty float y\nproperty float z\n") f.write("property uchar red\nproperty uchar green\nproperty uchar blue\n") f.write("end_header\n") for i in range(len(pts)): r, g, b = (np.clip(col[i], 0, 1) * 255).astype(np.uint8) f.write(f"{pts[i,0]:.6f} {pts[i,1]:.6f} {pts[i,2]:.6f} {r} {g} {b}\n") # ---------------------------------------------------------------- main def main(): ap = argparse.ArgumentParser() ap.add_argument("--transforms", required=True) ap.add_argument("--img_dir", required=True) ap.add_argument("--outdir", required=True) ap.add_argument("--steps", type=int, default=30000) ap.add_argument("--batch", type=int, default=4096) ap.add_argument("--lr", type=float, default=5e-4) ap.add_argument("--near", type=float, default=2.0) ap.add_argument("--far", type=float, default=12.0) ap.add_argument("--Nc", type=int, default=64) ap.add_argument("--Nf", type=int, default=64) ap.add_argument("--eval_every", type=int, default=2000) ap.add_argument("--export_only", action="store_true") ap.add_argument("--checkpoint", default=None) args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" print("device:", device, "| cuda:", torch.cuda.get_device_name(0) if device == "cuda" else "N/A") os.makedirs(args.outdir, exist_ok=True) imgs, c2ws, names, intr = load_dataset(args.transforms, args.img_dir) print(f"loaded {len(imgs)} images, res {intr[4]}x{intr[5]}, fx={intr[0]:.2f}") model = NeRF().to(device) opt = torch.optim.Adam(model.parameters(), lr=args.lr) start = 0 ckpt = args.checkpoint or os.path.join(args.outdir, "nerf_ckpt.pt") if os.path.exists(ckpt): d = torch.load(ckpt, map_location=device) model.load_state_dict(d["model"]); opt.load_state_dict(d["opt"]); start = d["step"] print(f"resumed from step {start}") if not args.export_only: t0 = time.time() for step in range(start, args.steps): model.train() rays_o, rays_d, cols = sample_batch(imgs, c2ws, intr, args.batch, device) (rgb_c, _, _), (rgb_f, _, _) = render_rays( model, rays_o, rays_d, args.near, args.far, args.Nc, args.Nf) loss = F.mse_loss(rgb_c, cols) + F.mse_loss(rgb_f, cols) opt.zero_grad(); loss.backward(); opt.step() if step % 100 == 0: print(f"step {step} loss {loss.item():.5f} ({100*(time.time()-t0)/max(1,step-start+1):.2f} s/100it)", flush=True) if (step + 1) % args.eval_every == 0: torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step + 1}, ckpt) # quick PSNR on a held-out view psnr = eval_psnr(model, imgs, c2ws, intr, device, args, test_idx=len(imgs)//2) print(f" [eval] PSNR on view {len(imgs)//2}: {psnr:.3f} dB", flush=True) torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": args.steps}, ckpt) print(f"training done in {(time.time()-t0)/60:.1f} min, saved {ckpt}") # export point cloud out_ply = os.path.join(args.outdir, "nerf_pointcloud.ply") print("exporting NeRF point cloud ...") export_pointcloud(model, imgs, c2ws, names, intr, out_ply, args.near, args.far, args.Nc, args.Nf, device, step=2, thresh=0.4) print("saved", out_ply) @torch.no_grad() def eval_psnr(model, imgs, c2ws, intr, device, args, test_idx): fx, fy, cx, cy, H, W = intr img = imgs[test_idx] c2w = torch.from_numpy(c2ws[test_idx]).to(device) us = torch.arange(W, device=device).float() vs = torch.arange(H, device=device).float() uv = torch.stack(torch.meshgrid(us, vs, indexing="ij"), dim=-1).reshape(-1, 2) model.eval() preds = [] for s in range(0, uv.shape[0], 4096): o, d = get_rays(c2w, fx, fy, cx, cy, H, W, uv[s:s+4096]) (_, _, _), (rgb, _, _) = render_rays(model, o, d, args.near, args.far, args.Nc, args.Nf, perturb=False) preds.append(rgb) pred = torch.cat(preds, 0).reshape(H, W, 3).cpu().numpy() gt = img mse = ((pred - gt) ** 2).mean() return -10 * np.log10(mse + 1e-9) if __name__ == "__main__": main()