""" BrepGaussian-adapted pipeline for CAD reconstruction from multi-view images. Adaptation of the BrepGaussian paper (arXiv 2602.21105, CVPR 2026) for the user's dataset of 4-5 CAD part views. Key adaptations: 1. POSE ESTIMATION: Uses MASt3R (3D foundation model) instead of COLMAP, because COLMAP fails on textureless CAD surfaces (verified). 2. RASTERIZER: Uses gsplat (pure PyTorch) instead of the custom diff-surfel-segment-rasterization CUDA extension (no nvcc available). 3. EDGE MAPS: Generated via OpenCV Canny edge detection on input images (BrepGaussian uses NEF_code rendering pipeline for synthetic data). 4. MASKS: Generated via simple region-based segmentation (BrepGaussian uses a fine-tuned SAM; we approximate with edge-guided watershed segmentation). 5. B-REP FITTING: Uses the official BrepGaussian BRep/ module unchanged (surface_fitting_pipeline.py) on the labeled point cloud. Pipeline stages: Stage 1: 3DGS training with photometric + edge supervision Stage 2: Instance feature learning with contrastive loss Export: Labeled point cloud (merged.pcd with label + edge fields) B-Rep: Analytic surface fitting (planes/cylinders) + mesh export Official repo: https://github.com/yjx2851/BrepGaussian """ import os import sys import json import glob import shutil import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import cv2 import trimesh from PIL import Image from tqdm import tqdm # Add BrepGaussian BRep module to path BREPGAUSSIAN_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "BrepGaussian") sys.path.insert(0, os.path.join(BREPGAUSSIAN_DIR, "BRep")) sys.path.insert(0, os.path.join(BREPGAUSSIAN_DIR, "GS")) PARTS_DIR = "data/parts" MAST3R_DIR = "data/mast3r_results" OUTPUT_DIR = "data/brepgaussian_results" DEVICE = "cuda" def load_poses_and_images(part): """Load MASt3R-estimated poses and images for a part.""" part_mast3r = os.path.join(MAST3R_DIR, part) colmap_dir = os.path.join(part_mast3r, "colmap_sparse") # Load transforms.json (NeRF convention, cam2world) with open(os.path.join(colmap_dir, "transforms.json")) as f: transforms = json.load(f) frames = transforms["frames"] images = [] poses = [] # cam2world intrinsics = [] for frame in frames: img_path = os.path.join(colmap_dir, frame["file_path"]) img = np.array(Image.open(img_path).convert("RGB")) h_actual, w_actual = img.shape[:2] # Use actual image dims (transforms.json w/h from MASt3R may differ slightly) images.append(img) c2w = np.array(frame["transform_matrix"]) poses.append(c2w) fx = frame["fl_x"] cx, cy = frame["cx"], frame["cy"] w, h = w_actual, h_actual # Scale intrinsics if MASt3R resolution differs from actual scale_x = w_actual / frame.get("w", w_actual) scale_y = h_actual / frame.get("h", h_actual) intrinsics.append({ "fx": fx * scale_x, "fy": frame["fl_y"] * scale_y, "cx": cx * scale_x, "cy": cy * scale_y, "w": w, "h": h }) return images, poses, intrinsics def generate_edge_maps(images): """Generate edge maps using Canny edge detection (Stage 1 supervision).""" edge_maps = [] for img in images: gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) # Canny with auto thresholds median = np.median(gray) lower = int(max(0, 0.66 * median)) upper = int(min(255, 1.33 * median)) edges = cv2.Canny(gray, lower, upper) # Dilate slightly for better supervision kernel = np.ones((3, 3), np.uint8) edges = cv2.dilate(edges, kernel, iterations=1) edge_maps.append(edges) return edge_maps def generate_instance_masks(images, edge_maps): """Generate per-face instance masks using edge-guided watershed segmentation. This approximates BrepGaussian's fine-tuned SAM with a simpler approach.""" masks = [] # list of per-image label maps (integer instance IDs) for img, edges in zip(images, edge_maps): gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) # Distance transform for watershed markers dist = cv2.distanceTransform(255 - edges, cv2.DIST_L2, 5) # Threshold to get sure foreground _, sure_fg = cv2.threshold(dist, 0.3 * dist.max(), 255, 0) sure_fg = sure_fg.astype(np.uint8) # Sure background (dilated) kernel = np.ones((5, 5), np.uint8) sure_bg = cv2.dilate(edges, kernel, iterations=3) sure_bg = 255 - sure_bg # Unknown region unknown = cv2.subtract(sure_bg, sure_fg) # Marker labelling num_labels, markers = cv2.connectedComponents(sure_fg) # Add 1 to all labels so background is 1 (not 0) markers = markers + 1 markers[unknown == 255] = 0 # Watershed img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) markers = cv2.watershed(img_bgr, markers) # markers: -1 = boundary, 1 = background, 2+ = instances markers[markers == -1] = 0 markers[markers == 1] = 0 # background masks.append(markers) # Find consistent instance labels across views by matching colors # (Simplified: assign unique global labels per view) global_masks = [] offset = 0 for m in masks: unique = np.unique(m) unique = unique[unique > 0] remap = np.zeros(m.max() + 1, dtype=np.int32) for u in unique: offset += 1 remap[u] = offset global_m = remap[m] global_masks.append(global_m) return global_masks def quaternions_to_rotmat(quats): """Convert quaternions (w,x,y,z) to rotation matrices.""" w, x, y, z = quats.unbind(-1) R = torch.stack([ 1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y), 2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x), 2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y), ], dim=-1).reshape(-1, 3, 3) return R def c2w_to_w2c(c2w): """Convert camera-to-world to world-to-camera (view matrix).""" R = c2w[:3, :3] t = c2w[:3, 3] # World-to-camera R_w2c = R.T t_w2c = -R_w2c @ t return R_w2c, t_w2c class BrepGaussianModel(nn.Module): """3DGS model with edge scalar + instance features (BrepGaussian adaptation). Uses gsplat for rasterization (pure PyTorch, no CUDA compilation needed). """ def __init__(self, n_gaussians=50000, sh_degree=3, feature_dim=16): super().__init__() self.n_gaussians = n_gaussians self.sh_degree = sh_degree self.feature_dim = feature_dim # Gaussian parameters self._xyz = nn.Parameter(torch.zeros(n_gaussians, 3)) self._features_dc = nn.Parameter(torch.zeros(n_gaussians, 1, 3)) self._features_rest = nn.Parameter(torch.zeros(n_gaussians, sh_degree ** 2, 3)) self._scaling = nn.Parameter(torch.ones(n_gaussians, 3) * 0.01) self._rotation = nn.Parameter(torch.zeros(n_gaussians, 4)) self._rotation.data[:, 0] = 1.0 # identity quaternion self._opacity = nn.Parameter(torch.zeros(n_gaussians, 1) - 5.0) # sigmoid -> ~0 # BrepGaussian-specific: per-Gaussian edge scalar self._edge = nn.Parameter(torch.zeros(n_gaussians, 1)) # BrepGaussian-specific: per-Gaussian instance feature (Stage 2) self._features = nn.Parameter(torch.zeros(n_gaussians, feature_dim)) def initialize_from_pointcloud(self, pcd_points, pcd_colors): """Initialize Gaussians from a point cloud.""" n = len(pcd_points) if n > self.n_gaussians: idx = np.random.RandomState(42).choice(n, self.n_gaussians, replace=False) pcd_points = pcd_points[idx] pcd_colors = pcd_colors[idx] elif n < self.n_gaussians: # Repeat to fill idx = np.random.RandomState(42).choice(n, self.n_gaussians, replace=True) pcd_points = pcd_points[idx] pcd_colors = pcd_colors[idx] with torch.no_grad(): self._xyz.copy_(torch.tensor(pcd_points, dtype=torch.float32, device=DEVICE)) # SH DC from colors colors = torch.tensor(pcd_colors, dtype=torch.float32, device=DEVICE) / 255.0 colors = (colors - 0.5) / 0.28209479177387814 # C0 self._features_dc.copy_(colors.unsqueeze(1)) # Scaling from point cloud density pts = torch.tensor(pcd_points, dtype=torch.float32, device=DEVICE) dists = torch.cdist(pts, pts) k = min(20, len(pts) - 1) knn_dists = dists.topk(k + 1, largest=False, dim=1)[0][:, 1:] avg_dist = knn_dists.mean(dim=1) scale = avg_dist.unsqueeze(1).repeat(1, 3) self._scaling.copy_(torch.log(scale.clamp(min=0.001))) # Opacity self._opacity.data.fill_(0.0) def get_params(self, stage="stage1"): """Get parameters for optimization.""" params = [ {"params": [self._xyz], "lr": 0.00016, "name": "xyz"}, {"params": [self._features_dc], "lr": 0.0025, "name": "f_dc"}, {"params": [self._features_rest], "lr": 0.0025 / 20, "name": "f_rest"}, {"params": [self._scaling], "lr": 0.005, "name": "scaling"}, {"params": [self._rotation], "lr": 0.001, "name": "rotation"}, {"params": [self._opacity], "lr": 0.05, "name": "opacity"}, {"params": [self._edge], "lr": 0.005, "name": "edge"}, ] if stage == "stage2": params.append({"params": [self._features], "lr": 0.01, "name": "features"}) # Freeze geometry in stage 2 for p in params: if p["name"] in ["xyz", "f_dc", "f_rest", "scaling", "rotation"]: p["params"][0].requires_grad_(False) return params def render(self, viewmat, intrinsics, h, w, render_features=False): """Render the Gaussians from a viewpoint using pure-PyTorch renderer.""" from simple_renderer import rasterize_gaussians means3D = self._xyz # Opacity opacity = torch.sigmoid(self._opacity).squeeze(-1) # Scaling and rotation scales = torch.exp(self._scaling) rotations = F.normalize(self._rotation, dim=-1) # View matrix (world-to-camera) as 4x4 viewmat = viewmat.to(DEVICE).float() K = torch.tensor([ [intrinsics["fx"], 0, intrinsics["cx"]], [0, intrinsics["fy"], intrinsics["cy"]], [0, 0, 1] ], dtype=torch.float32, device=DEVICE) if render_features: # Render RGB + edge + features edge = torch.sigmoid(self._edge) # [N, 1] feats = torch.sigmoid(self._features) # [N, 16] colors = torch.cat([ torch.clamp(0.5 + 0.28209479177387814 * self._features_dc.squeeze(1), 0, 1), edge, feats ], dim=1) # [N, 20] render = rasterize_gaussians( means3D, rotations, scales, opacity, colors, viewmat, K, h, w ) return render else: # Render RGB + edge only edge = torch.sigmoid(self._edge) colors = torch.cat([ torch.clamp(0.5 + 0.28209479177387814 * self._features_dc.squeeze(1), 0, 1), edge ], dim=1) # [N, 4] render = rasterize_gaussians( means3D, rotations, scales, opacity, colors, viewmat, K, h, w ) return render def train_stage1(model, images, poses, intrinsics, edge_maps, n_iters=7000): """Stage 1: Photometric + edge supervision training.""" print(f"\n Stage 1: {n_iters} iterations (photometric + edge)") optimizer = torch.optim.Adam(model.get_params("stage1")) # Prepare data tensors - resize all to consistent size if needed imgs_t = [] edges_t = [] for im, em in zip(images, edge_maps): h, w = im.shape[:2] imgs_t.append(torch.tensor(im / 255.0, dtype=torch.float32, device=DEVICE)) edges_t.append(torch.tensor(em / 255.0, dtype=torch.float32, device=DEVICE)) losses = [] for i in tqdm(range(n_iters), desc="Stage1"): idx = i % len(images) c2w = poses[idx] R_w2c, t_w2c = c2w_to_w2c(c2w) viewmat = torch.eye(4, device=DEVICE) viewmat[:3, :3] = torch.tensor(R_w2c, device=DEVICE, dtype=torch.float32) viewmat[:3, 3] = torch.tensor(t_w2c, device=DEVICE, dtype=torch.float32) intr = intrinsics[idx] h, w = intr["h"], intr["w"] try: render = model.render(viewmat, intr, h, w, render_features=False) except Exception as e: if i == 0: print(f" Render error: {e}") continue rgb_render = render[:, :, :3] edge_render = render[:, :, 3:4] gt_rgb = imgs_t[idx][:h, :w] gt_edge = edges_t[idx][:h, :w].unsqueeze(-1) # Photometric loss (L1 + SSIM) l1 = F.l1_loss(rgb_render, gt_rgb) # Simplified SSIM loss ssim_loss = 1.0 - simple_ssim(rgb_render, gt_rgb) lambda_dssim = 0.2 photo_loss = (1 - lambda_dssim) * l1 + lambda_dssim * ssim_loss # Edge loss edge_loss = F.l1_loss(edge_render, gt_edge) loss = photo_loss + 0.1 * edge_loss optimizer.zero_grad() loss.backward() optimizer.step() if i % 500 == 0: print(f" iter {i}: photo={photo_loss.item():.4f}, edge={edge_loss.item():.4f}, total={loss.item():.4f}") losses.append({"iter": i, "photo": photo_loss.item(), "edge": edge_loss.item(), "total": loss.item()}) return losses def simple_ssim(pred, gt, window_size=11): """Simplified SSIM computation.""" if pred.shape[-1] != 1: pred = pred.mean(-1, keepdim=True) gt = gt.mean(-1, keepdim=True) C1, C2 = 0.01 ** 2, 0.03 ** 2 mu_p = F.avg_pool2d(pred.permute(2, 0, 1).unsqueeze(0), window_size, stride=1, padding=window_size // 2) mu_g = F.avg_pool2d(gt.permute(2, 0, 1).unsqueeze(0), window_size, stride=1, padding=window_size // 2) mu_p, mu_g = mu_p.squeeze(0).permute(1, 2, 0), mu_g.squeeze(0).permute(1, 2, 0) sigma_p = F.avg_pool2d((pred - mu_p).permute(2, 0, 1).unsqueeze(0), window_size, stride=1, padding=window_size // 2).squeeze(0).permute(1, 2, 0) sigma_g = F.avg_pool2d((gt - mu_g).permute(2, 0, 1).unsqueeze(0), window_size, stride=1, padding=window_size // 2).squeeze(0).permute(1, 2, 0) sigma_pg = F.avg_pool2d(((pred - mu_p) * (gt - mu_g)).permute(2, 0, 1).unsqueeze(0), window_size, stride=1, padding=window_size // 2).squeeze(0).permute(1, 2, 0) ssim = ((2 * mu_p * mu_g + C1) * (2 * sigma_pg + C2)) / ((mu_p ** 2 + mu_g ** 2 + C1) * (sigma_p + sigma_g + C2)) return ssim.mean() def train_stage2(model, images, poses, intrinsics, masks, n_iters=3000): """Stage 2: Instance feature learning with contrastive loss.""" print(f"\n Stage 2: {n_iters} iterations (instance features)") optimizer = torch.optim.Adam(model.get_params("stage2")) masks_t = [torch.tensor(m, dtype=torch.long, device=DEVICE) for m in masks] losses = [] for i in tqdm(range(n_iters), desc="Stage2"): idx = i % len(images) c2w = poses[idx] R_w2c, t_w2c = c2w_to_w2c(c2w) viewmat = torch.eye(4, device=DEVICE) viewmat[:3, :3] = torch.tensor(R_w2c, device=DEVICE, dtype=torch.float32) viewmat[:3, 3] = torch.tensor(t_w2c, device=DEVICE, dtype=torch.float32) intr = intrinsics[idx] h, w = intr["h"], intr["w"] try: render = model.render(viewmat, intr, h, w, render_features=True) except Exception as e: if i == 0: print(f" Render error: {e}") continue # render: [H, W, 20] = RGB(3) + edge(1) + features(16) feat_render = render[:, :, 4:] # [H, W, 16] gt_mask = masks_t[idx][:h, :w] # Contrastive loss: pull same-instance features together, push different apart # For each instance label, compute feature centroid unique_labels = torch.unique(gt_mask) unique_labels = unique_labels[unique_labels > 0] if len(unique_labels) < 2: continue loss = torch.tensor(0.0, device=DEVICE, requires_grad=True) centroids = {} for label in unique_labels: mask = (gt_mask == label) if mask.sum() < 10: continue feats = feat_render[mask] # [K, 16] centroid = feats.mean(0) centroids[label.item()] = centroid # Pull features toward centroid loss = loss + F.mse_loss(feats, centroid.unsqueeze(0).expand_as(feats)) # Push centroids apart (maximize distance between different instances) labels_list = list(centroids.keys()) if len(labels_list) >= 2: for j in range(len(labels_list)): for k in range(j + 1, len(labels_list)): c1, c2 = centroids[labels_list[j]], centroids[labels_list[k]] # Maximize cosine distance -> minimize cosine similarity cos_sim = F.cosine_similarity(c1.unsqueeze(0), c2.unsqueeze(0)) loss = loss + 0.5 * cos_sim if loss.requires_grad: optimizer.zero_grad() loss.backward() optimizer.step() if i % 500 == 0: print(f" iter {i}: contrastive_loss={loss.item():.4f}") losses.append({"iter": i, "contrastive": loss.item()}) return losses def export_labeled_pcd(model, output_path): """Export labeled point cloud in BrepGaussian's merged.pcd format. Fields: x y z label edge - label: instance label from feature clustering - edge: edge scalar value """ with torch.no_grad(): xyz = model._xyz.cpu().numpy() edge_vals = torch.sigmoid(model._edge).cpu().numpy().squeeze() feats = torch.sigmoid(model._features).cpu().numpy() # Cluster features to assign instance labels from sklearn.cluster import KMeans n_clusters = min(10, max(2, len(xyz) // 5000)) kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10) labels = kmeans.fit_predict(feats) # Threshold edge values edge_threshold = np.median(edge_vals) + 0.5 * np.std(edge_vals) edge_binary = (edge_vals > edge_threshold).astype(float) # Write PCD with open(output_path, "w") as f: f.write("# .PCD v0.7 - Point Cloud Data file\n") f.write("VERSION 0.7\n") f.write("FIELDS x y z label edge\n") f.write("SIZE 4 4 4 4 4\n") f.write("TYPE F F F I F\n") f.write("COUNT 1 1 1 1 1\n") f.write(f"WIDTH {len(xyz)}\n") f.write("HEIGHT 1\n") f.write("VIEWPOINT 0 0 0 1 0 0 0\n") f.write(f"POINTS {len(xyz)}\n") f.write("DATA ascii\n") for i in range(len(xyz)): f.write(f"{xyz[i,0]:.6f} {xyz[i,1]:.6f} {xyz[i,2]:.6f} {int(labels[i])} {edge_binary[i]:.6f}\n") print(f" Exported labeled PCD: {output_path} ({len(xyz)} points, {n_clusters} clusters)") return output_path def run_brep_fitting(pcd_path, output_dir): """Run the official BrepGaussian B-Rep fitting pipeline.""" print(f"\n B-Rep fitting...") os.makedirs(output_dir, exist_ok=True) # Import the official pipeline from surface_fitting_pipeline import run_pipeline run_pipeline(pcd_path=pcd_path, output_dir=output_dir, visualize=False) print(f" B-Rep output: {output_dir}") return output_dir def process_part(part): """Full BrepGaussian pipeline for one part.""" print(f"\n{'='*60}") print(f"BrepGaussian pipeline: {part}") print(f"{'='*60}") out_dir = os.path.join(OUTPUT_DIR, part) os.makedirs(out_dir, exist_ok=True) # 1. Load poses + images from MASt3R print(" [1/6] Loading poses and images...") images, poses, intrinsics = load_poses_and_images(part) print(f" {len(images)} images, sizes: {[im.shape for im in images]}") # 2. Generate edge maps + masks print(" [2/6] Generating edge maps and instance masks...") edge_maps = generate_edge_maps(images) masks = generate_instance_masks(images, edge_maps) # Save edge maps and masks for i, (em, m) in enumerate(zip(edge_maps, masks)): cv2.imwrite(os.path.join(out_dir, f"edge_{i:02d}.png"), em) # Save mask as colored mask_colored = np.zeros((m.shape[0], m.shape[1], 3), dtype=np.uint8) for label in np.unique(m): if label > 0: color = np.random.RandomState(int(label)).randint(0, 255, 3) mask_colored[m == label] = color cv2.imwrite(os.path.join(out_dir, f"mask_{i:02d}.png"), mask_colored) # 3. Load MASt3R point cloud for initialization print(" [3/6] Loading MASt3R point cloud for initialization...") import open3d as o3d pcd_path = os.path.join(MAST3R_DIR, part, "pointcloud.ply") pcd = o3d.io.read_point_cloud(pcd_path) pts = np.asarray(pcd.points) cols = (np.asarray(pcd.colors) * 255).astype(np.uint8) if pcd.has_colors() else np.tile([128, 128, 128], (len(pts), 1)) print(f" {len(pts)} points for initialization") # 4. Initialize and train Stage 1 print(" [4/6] Stage 1 training (geometry + edges)...") n_gaussians = min(10000, len(pts)) model = BrepGaussianModel(n_gaussians=n_gaussians, sh_degree=3, feature_dim=16) model.to(DEVICE) model.initialize_from_pointcloud(pts, cols) stage1_iters = 2000 stage1_losses = train_stage1(model, images, poses, intrinsics, edge_maps, n_iters=stage1_iters) # Save stage 1 checkpoint torch.save({ "xyz": model._xyz, "features_dc": model._features_dc, "features_rest": model._features_rest, "scaling": model._scaling, "rotation": model._rotation, "opacity": model._opacity, "edge": model._edge, }, os.path.join(out_dir, "stage1_checkpoint.pth")) # 5. Stage 2 training print(" [5/6] Stage 2 training (instance features)...") stage2_iters = 1000 stage2_losses = train_stage2(model, images, poses, intrinsics, masks, n_iters=stage2_iters) # Export labeled point cloud pcd_out = os.path.join(out_dir, "merged.pcd") export_labeled_pcd(model, pcd_out) # 6. B-Rep fitting print(" [6/6] B-Rep surface fitting...") brep_dir = os.path.join(out_dir, "brep_output") try: run_brep_fitting(pcd_out, brep_dir) except Exception as e: import traceback traceback.print_exc() print(f" B-Rep fitting failed: {e}") brep_dir = None # Save Gaussian splat model as PLY with torch.no_grad(): xyz = model._xyz.cpu().numpy() cols = (torch.clamp(0.5 + 0.28209479177387814 * model._features_dc.squeeze(1), 0, 1).cpu().numpy() * 255).astype(np.uint8) gs_ply = trimesh.PointCloud(xyz, colors=cols) gs_ply.export(os.path.join(out_dir, "gaussian_splat.ply")) # Save metadata meta = { "part": part, "method": "BrepGaussian-adapted (simple PyTorch renderer + MASt3R poses + OpenCV edges + watershed masks + official B-Rep fitting)", "n_images": len(images), "n_gaussians": n_gaussians, "stage1_iters": stage1_iters, "stage2_iters": stage2_iters, "stage1_losses": stage1_losses, "stage2_losses": stage2_losses, } with open(os.path.join(out_dir, "metadata.json"), "w") as f: json.dump(meta, f, indent=2) return out_dir def main(): os.makedirs(OUTPUT_DIR, exist_ok=True) results = {} parts = sorted([d for d in os.listdir(PARTS_DIR) if os.path.isdir(os.path.join(PARTS_DIR, d, "input"))]) for part in parts: # Skip if MASt3R results don't exist if not os.path.exists(os.path.join(MAST3R_DIR, part)): print(f"Skipping {part}: no MASt3R results") continue try: out_dir = process_part(part) results[part] = ("OK", out_dir) except Exception as e: import traceback traceback.print_exc() results[part] = ("FAILED", str(e)) print(f"\n{'='*60}") print("BREPGAUSSIAN PIPELINE SUMMARY") print(f"{'='*60}") for part, (status, msg) in results.items(): print(f" {part}: {status} - {msg}") return results if __name__ == "__main__": main()