# -*- coding: utf-8 -*- """ Plane-sweep multi-view stereo (own implementation, PyTorch/CUDA). Given N calibrated views (COLMAP-estimated poses) this reconstructs a dense coloured point cloud: for each reference view: 1. pick K nearest source views (by camera-centre distance) 2. sweep D fronto-parallel depth planes (uniform in inverse depth) 3. warp every source onto each plane, score photometric consistency (SSD patch) 4. best depth per pixel = argmax score; sub-pixel refine 5. keep pixels above a consistency threshold, unproject to 3D fuse all views -> voxel downsample -> write PLY Usage: python mvs_planesweep.py --poses output/colmap/cameras_export/poses_c2w.npz \ --img_dir mvs_project/images --outdir output/mvs --D 192 --K 10 """ import os, sys, argparse import numpy as np import torch import torch.nn.functional as F def load_images(img_dir, names): from PIL import Image imgs_rgb, imgs_gray = [], [] for n in names: p = os.path.join(img_dir, n) rgb = np.array(Image.open(p).convert("RGB")).astype(np.float32) / 255.0 gray = rgb.mean(axis=-1) imgs_rgb.append(rgb); imgs_gray.append(gray) return np.stack(imgs_rgb), np.stack(imgs_gray) def build_ref_grid(H, W, device): ys, xs = torch.meshgrid(torch.arange(H, device=device), torch.arange(W, device=device), indexing="ij") ones = torch.ones_like(xs) return torch.stack([xs.float(), ys.float(), ones], dim=-1) # [H,W,3] @torch.no_grad() def plane_sweep_view(ref_gray, ref_rgb, src_gray, src_rgb, K, c2w_ref, c2w_srcs, D, d_min, d_max, device): """One reference view -> (depth_map, score_map).""" H, W = ref_gray.shape K_t = torch.tensor(K, dtype=torch.float32, device=device) Kinv = torch.inverse(K_t) c2w_ref = torch.tensor(c2w_ref, dtype=torch.float32, device=device) R_ref, t_ref = c2w_ref[:3, :3], c2w_ref[:3, 3] ref_t = torch.tensor(ref_gray, dtype=torch.float32, device=device)[None, None] srcs = [torch.tensor(s, dtype=torch.float32, device=device)[None, None] for s in src_gray] src_poses = [torch.tensor(c, dtype=torch.float32, device=device) for c in c2w_srcs] # depth planes (inverse-depth uniform) inv_d = torch.linspace(1.0 / d_max, 1.0 / d_min, D, device=device) depths = 1.0 / inv_d # [D] grid = build_ref_grid(H, W, device) # [H,W,3] pixel coords # aggregate photo-consistency: sum of SSD across sources per plane acc = torch.zeros(D, H, W, device=device) # [D,H,W] n = torch.tensor([0.0, 0.0, 1.0], device=device) for s, c_src in zip(srcs, src_poses): R_src, t_src = c_src[:3, :3], c_src[:3, 3] # X_src = R X_ref + t (ref frame -> src frame) R = R_src.T @ R_ref t = R_src.T @ (t_ref - t_src) for di, d in enumerate(depths): # H_ref2src = K @ (R + t n^T / d) @ K^-1 (n=[0,0,1] in ref frame) H3d = R + (t.unsqueeze(1) @ n.unsqueeze(0)) / d # [3,3] H_mat = K_t @ H3d @ Kinv # ref -> src pixels p_src = grid @ H_mat.T # [H,W,3] u = p_src[..., 0] / p_src[..., 2] v = p_src[..., 1] / p_src[..., 2] su = 2.0 * u / (W - 1) - 1.0 sv = 2.0 * v / (H - 1) - 1.0 sample_grid = torch.stack([su, sv], dim=-1)[None] # [1,H,W,2] warped = F.grid_sample(s, sample_grid, align_corners=True, padding_mode="border") diff = (warped[0, 0] - ref_t[0, 0]) patch = F.avg_pool2d((diff ** 2)[None, None], kernel_size=5, stride=1, padding=2)[0, 0] acc[di] += patch / len(srcs) # average SSD over sources # best depth = argmin average SSD over planes best_idx = acc.argmin(dim=0) # [H,W] best_depth = depths[best_idx] best_cost = acc.gather(0, best_idx[None])[0] # [H,W] min avg SSD best_score = -best_cost # higher = better return best_depth.cpu().numpy(), best_score.cpu().numpy() def main(): ap = argparse.ArgumentParser() ap.add_argument("--poses", required=True) ap.add_argument("--img_dir", required=True) ap.add_argument("--outdir", required=True) ap.add_argument("--D", type=int, default=192) ap.add_argument("--K", type=int, default=10) ap.add_argument("--d_min", type=float, default=1.0) ap.add_argument("--d_max", type=float, default=10.0) ap.add_argument("--score_thresh", type=float, default=-0.02) ap.add_argument("--tex_thresh", type=float, default=0.01) ap.add_argument("--max_views", type=int, default=0) args = ap.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" print("device:", device) os.makedirs(args.outdir, exist_ok=True) d = np.load(args.poses, allow_pickle=True) names = list(d["names"]); c2w = d["c2w"] fx, fy, cx, cy = float(d["fx"]), float(d["fy"]), float(d["cx"]), float(d["cy"]) K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], np.float64) H = W = int(d["w"]) if args.max_views: names, c2w = names[:args.max_views], c2w[:args.max_views] N = len(names) rgb, gray = load_images(args.img_dir, names) print(f"loaded {N} images {W}x{H}, fx={fx:.2f}") # local texture map (std over 5x5) to reject textureless sky/surfaces tex = [] for g in gray: g = torch.tensor(g, dtype=torch.float32, device=device)[None, None] mu = F.avg_pool2d(g, 5, stride=1, padding=2) var = F.avg_pool2d(g ** 2, 5, stride=1, padding=2) - mu ** 2 tex.append(torch.sqrt(var.clamp(min=0))[0, 0].cpu().numpy()) tex = np.stack(tex) # camera centres for neighbour selection centers = c2w[:, :3, 3] all_pts, all_col = [], [] for i in range(N): # select K nearest sources (by centre distance) dist = np.linalg.norm(centers - centers[i], axis=1) order = np.argsort(dist) src_idx = [j for j in order if j != i][:args.K] depth, score = plane_sweep_view(gray[i], rgb[i], [gray[j] for j in src_idx], [rgb[j] for j in src_idx], K, c2w[i], [c2w[j] for j in src_idx], args.D, args.d_min, args.d_max, device) mask = (score > args.score_thresh) & (tex[i] > args.tex_thresh) # unproject ys, xs = np.nonzero(mask) z = depth[ys, xs] Xc = (xs - cx) * z / fx Yc = (ys - cy) * z / fy pts_cam = np.stack([Xc, Yc, z], axis=-1) # [M,3] camera coords R = c2w[i][:3, :3]; t = c2w[i][:3, 3] pts_world = (R @ pts_cam.T).T + t col = rgb[i][ys, xs] # [M,3] 0..1 all_pts.append(pts_world); all_col.append(col) print(f"view {i+1}/{N} ({names[i]}): {len(ys)} points", flush=True) P = np.concatenate(all_pts, 0); C = np.concatenate(all_col, 0) print(f"total raw points: {len(P)}") # voxel downsample 1cm keys = np.floor(P / 0.01).astype(np.int64) _, idx = np.unique(keys, axis=0, return_index=True) P, C = P[idx], C[idx] print(f"after 1cm voxel downsample: {len(P)}") out = os.path.join(args.outdir, "mvs_pointcloud.ply") write_ply(out, P, C) print("saved", out) def write_ply(path, pts, col): 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") if __name__ == "__main__": main()