| """Edge existence classifier for S23DR 2026. |
| |
| Given merged 3D vertices, predict which vertex pairs form real wireframe edges. |
| Enables cross-view edge prediction that gestalt-adjacency cannot provide. |
| """ |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
|
|
| class EdgeExistenceModel(nn.Module): |
| def __init__(self, input_dim=32, hidden_dim=64): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(input_dim, hidden_dim), |
| nn.BatchNorm1d(hidden_dim), |
| nn.ReLU(), |
| nn.Dropout(0.2), |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Linear(hidden_dim, 1), |
| ) |
|
|
| def forward(self, x): |
| return self.net(x).squeeze(-1) |
|
|
|
|
| class EdgeExistenceTrainer: |
| def __init__(self, device='cuda', lr=1e-3, pos_weight=10.0): |
| self.device = device |
| self.model = EdgeExistenceModel().to(device) |
| self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=1e-4) |
| pos_weight_tensor = torch.tensor([pos_weight], device=device) |
| self.criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight_tensor) |
|
|
| def train_step(self, features, labels): |
| self.model.train() |
| self.optimizer.zero_grad() |
| logits = self.model(features.to(self.device)) |
| loss = self.criterion(logits, labels.float().to(self.device)) |
| loss.backward() |
| self.optimizer.step() |
| return loss.item() |
|
|
| @torch.no_grad() |
| def predict_proba(self, features): |
| self.model.eval() |
| if isinstance(features, np.ndarray): |
| features = torch.from_numpy(features).float() |
| return torch.sigmoid(self.model(features.to(self.device))).cpu().numpy() |
|
|
| def predict(self, features, threshold=0.4): |
| return self.predict_proba(features) >= threshold |
|
|
| def save(self, path): |
| import os |
| os.makedirs(os.path.dirname(path), exist_ok=True) |
| torch.save({'model_state_dict': self.model.state_dict()}, path) |
|
|
| def load(self, path): |
| ckpt = torch.load(path, map_location=self.device, weights_only=True) |
| self.model.load_state_dict(ckpt['model_state_dict']) |
|
|
|
|
| def generate_candidate_pairs(vertices, dist_thresh=8.0): |
| """Return all vertex index pairs within dist_thresh metres.""" |
| if len(vertices) < 2: |
| return [] |
| from scipy.spatial import cKDTree |
| tree = cKDTree(vertices) |
| pairs = list(tree.query_pairs(dist_thresh)) |
| return pairs |
|
|
|
|
| def label_candidate_pairs(pred_vertices, candidate_pairs, gt_vertices, gt_edges, match_th=0.5): |
| """Label each candidate pair 1 (positive) or 0 (negative). |
| |
| Positive: midpoint within match_th of a GT edge midpoint AND both endpoints |
| within match_th of a GT vertex. |
| """ |
| if len(gt_vertices) == 0 or len(gt_edges) == 0 or len(candidate_pairs) == 0: |
| return np.zeros(len(candidate_pairs), dtype=np.float32) |
|
|
| from scipy.spatial import cKDTree |
| gt_v = np.array(gt_vertices) |
| gt_tree = cKDTree(gt_v) |
|
|
| valid_gt_edges = [(a, b) for a, b in gt_edges if a < len(gt_v) and b < len(gt_v)] |
| if not valid_gt_edges: |
| return np.zeros(len(candidate_pairs), dtype=np.float32) |
|
|
| gt_mids = np.array([(gt_v[a] + gt_v[b]) / 2 for a, b in valid_gt_edges]) |
| mid_tree = cKDTree(gt_mids) |
|
|
| labels = np.zeros(len(candidate_pairs), dtype=np.float32) |
| for i, (a, b) in enumerate(candidate_pairs): |
| if a >= len(pred_vertices) or b >= len(pred_vertices): |
| continue |
| midpoint = (pred_vertices[a] + pred_vertices[b]) / 2 |
| mid_dist, _ = mid_tree.query(midpoint) |
| if mid_dist > match_th: |
| continue |
| da, _ = gt_tree.query(pred_vertices[a]) |
| db, _ = gt_tree.query(pred_vertices[b]) |
| if da < match_th and db < match_th: |
| labels[i] = 1.0 |
| return labels |
|
|
