"""Edge semantic classifier for S23DR 2026. Given a 3D wireframe (vertices + edges), classify each edge into one of 10 semantic types. Uses features from: gestalt segmentation, geometric properties, COLMAP point density. """ import torch import torch.nn as nn import torch.nn.functional as F import numpy as np from typing import Dict, List, Tuple, Optional try: from hoho2025.color_mappings import EDGE_CLASSES, gestalt_color_mapping, edge_color_mapping except ImportError: EDGE_CLASSES = {'cornice_return': 0, 'cornice_strip': 1, 'eave': 2, 'flashing': 3, 'hip': 4, 'rake': 5, 'ridge': 6, 'step_flashing': 7, 'transition_line': 8, 'valley': 9} NUM_EDGE_CLASSES = len(EDGE_CLASSES) def compute_edge_features(vertices, edges, colmap_rec=None, gestalt_images=None, camera_params=None): if len(edges) == 0: return np.zeros((0, 32)) degree = np.zeros(len(vertices), dtype=int) for a, b in edges: if a < len(vertices): degree[a] += 1 if b < len(vertices): degree[b] += 1 v_min = vertices.min(axis=0) if len(vertices) > 0 else np.zeros(3) v_max = vertices.max(axis=0) if len(vertices) > 0 else np.ones(3) v_range = v_max - v_min; v_range[v_range < 1e-6] = 1.0 building_diag = np.linalg.norm(v_range) if np.linalg.norm(v_range) > 1e-6 else 1.0 building_height = v_range[2] if v_range[2] > 1e-6 else 1.0 # Build adjacency for neighbor-context features adj = {i: [] for i in range(len(vertices))} for a, b in edges: if a < len(vertices) and b < len(vertices): adj[a].append(b); adj[b].append(a) sfm_points = None if colmap_rec is not None: try: pts = [p3d.xyz for pid, p3d in colmap_rec.points3D.items()] if pts: sfm_points = np.array(pts) except Exception: pass features = [] for a, b in edges: if a >= len(vertices) or b >= len(vertices): features.append(np.zeros(32)); continue v1, v2 = vertices[a], vertices[b] direction = v2 - v1; length = np.linalg.norm(direction) direction_norm = direction / length if length > 1e-6 else np.zeros(3) midpoint = (v1 + v2) / 2 vertical = np.array([0, 0, 1]) angle_with_vertical = np.arccos(np.clip(abs(np.dot(direction_norm, vertical)), 0, 1)) angle_with_horizontal = np.pi / 2 - angle_with_vertical rel_height_mid = (midpoint[2] - v_min[2]) / v_range[2] rel_height_v1 = (v1[2] - v_min[2]) / v_range[2] rel_height_v2 = (v2[2] - v_min[2]) / v_range[2] delta_z_norm = abs(v2[2] - v1[2]) / building_height cos_x, cos_y, cos_z = abs(direction_norm[0]), abs(direction_norm[1]), abs(direction_norm[2]) length_norm = length / building_diag dxy = np.sqrt((v2[0]-v1[0])**2 + (v2[1]-v1[1])**2) slope = np.arctan2(abs(v2[2]-v1[2]), max(dxy, 1e-6)) is_near_top = float(rel_height_mid > 0.75) is_near_bottom = float(rel_height_mid < 0.25) is_horizontal = float(angle_with_horizontal < np.deg2rad(15)) is_diagonal = float(np.deg2rad(15) <= angle_with_horizontal < np.deg2rad(60)) max_degree = max(degree[a], degree[b]) min_degree = min(degree[a], degree[b]) # Neighbor height context nbrs_a = [vertices[n][2] for n in adj[a] if n != b and n < len(vertices)] nbrs_b = [vertices[n][2] for n in adj[b] if n != a and n < len(vertices)] nbr_mean_z_a = np.mean(nbrs_a) if nbrs_a else midpoint[2] nbr_mean_z_b = np.mean(nbrs_b) if nbrs_b else midpoint[2] nbr_rel_a = (nbr_mean_z_a - v_min[2]) / v_range[2] nbr_rel_b = (nbr_mean_z_b - v_min[2]) / v_range[2] # Is this edge above or below neighboring edges? above_nbrs = float(rel_height_mid > max(nbr_rel_a, nbr_rel_b) - 0.05) density = 0.0 if sfm_points is not None and len(sfm_points) > 0: dists = np.linalg.norm(sfm_points - midpoint, axis=1) density = (dists < 1.0).sum() / max(len(sfm_points), 1) feat = np.array([ length_norm, # 0 direction_norm[0], # 1 direction_norm[1], # 2 direction_norm[2], # 3 angle_with_vertical, # 4 angle_with_horizontal, # 5 slope, # 6 rel_height_mid, # 7 rel_height_v1, # 8 rel_height_v2, # 9 delta_z_norm, # 10 cos_x, cos_y, cos_z, # 11-13 float(degree[a]), # 14 float(degree[b]), # 15 float(max_degree), # 16 float(min_degree), # 17 density, # 18 is_near_top, # 19 is_near_bottom, # 20 is_horizontal, # 21 is_diagonal, # 22 nbr_rel_a, # 23 nbr_rel_b, # 24 above_nbrs, # 25 float(len(adj[a])), # 26 num neighbors of a float(len(adj[b])), # 27 num neighbors of b float(rel_height_v1 > 0.7), # 28 v1 near top float(rel_height_v2 > 0.7), # 29 v2 near top float(rel_height_v1 < 0.3), # 30 v1 near bottom float(rel_height_v2 < 0.3), # 31 v2 near bottom ]) features.append(feat) return np.array(features) class EdgeClassifierMLP(nn.Module): def __init__(self, input_dim=32, hidden_dim=128, num_classes=NUM_EDGE_CLASSES, dropout=0.3): super().__init__() self.net = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim // 2), nn.ReLU(), nn.Linear(hidden_dim // 2, num_classes)) def forward(self, x): return self.net(x) class EdgeClassifierGNN(nn.Module): def __init__(self, node_dim=6, edge_feat_dim=32, hidden_dim=256, num_classes=NUM_EDGE_CLASSES, num_layers=3): super().__init__() self.node_encoder = nn.Sequential(nn.Linear(node_dim, hidden_dim), nn.ReLU()) self.edge_encoder = nn.Sequential(nn.Linear(edge_feat_dim, hidden_dim), nn.ReLU()) self.mp_layers = nn.ModuleList([MessagePassingLayer(hidden_dim) for _ in range(num_layers)]) self.classifier = nn.Sequential(nn.Linear(hidden_dim * 3, hidden_dim), nn.ReLU(), nn.Dropout(0.2), nn.Linear(hidden_dim, num_classes)) def forward(self, node_features, edge_index, edge_features): h = self.node_encoder(node_features); e = self.edge_encoder(edge_features) for mp in self.mp_layers: h = mp(h, edge_index, e) src_h = h[edge_index[0]]; dst_h = h[edge_index[1]] return self.classifier(torch.cat([src_h, dst_h, e], dim=1)) class MessagePassingLayer(nn.Module): def __init__(self, hidden_dim): super().__init__() self.message_fn = nn.Sequential(nn.Linear(hidden_dim * 3, hidden_dim), nn.ReLU()) self.update_fn = nn.Sequential(nn.Linear(hidden_dim * 2, hidden_dim), nn.ReLU()) def forward(self, h, edge_index, e): src, dst = edge_index messages = self.message_fn(torch.cat([h[src], h[dst], e], dim=1)) agg = torch.zeros_like(h); agg.scatter_add_(0, dst.unsqueeze(1).expand_as(messages), messages) return self.update_fn(torch.cat([h, agg], dim=1)) + h def prepare_gnn_input(vertices, edges, edge_features): if len(vertices) == 0: return None v_min = vertices.min(axis=0); v_range = vertices.max(axis=0) - v_min; v_range[v_range < 1e-6] = 1.0 node_features = np.concatenate([vertices, (vertices - v_min) / v_range], axis=1) if edges: src = [a for a, b in edges] + [b for a, b in edges]; dst = [b for a, b in edges] + [a for a, b in edges] edge_index = np.array([src, dst]); edge_features_both = np.concatenate([edge_features, edge_features], axis=0) else: edge_index = np.zeros((2, 0), dtype=int); edge_features_both = np.zeros((0, edge_features.shape[1] if len(edge_features) > 0 else 32)) return {'node_features': torch.from_numpy(node_features).float(), 'edge_index': torch.from_numpy(edge_index).long(), 'edge_features': torch.from_numpy(edge_features_both).float()} class EdgeClassifierTrainer: def __init__(self, model_type='mlp', device='cuda', lr=1e-3): self.device = device self.model = (EdgeClassifierMLP() if model_type == 'mlp' else EdgeClassifierGNN()).to(device) self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=1e-4) self.criterion = nn.CrossEntropyLoss(label_smoothing=0.1) def train_step_mlp(self, features, labels): self.model.train(); self.optimizer.zero_grad() logits = self.model(features.to(self.device)); loss = self.criterion(logits, labels.to(self.device)) loss.backward(); self.optimizer.step(); return loss.item() def train_step_gnn(self, vertices_np, edges, features_np, labels_tensor): self.model.train(); self.optimizer.zero_grad() gnn_input = prepare_gnn_input(vertices_np, edges, features_np) if gnn_input is None or gnn_input['edge_index'].shape[1] == 0: return 0.0 logits = self.model(gnn_input['node_features'].to(self.device), gnn_input['edge_index'].to(self.device), gnn_input['edge_features'].to(self.device)) loss = self.criterion(logits[:len(edges)], labels_tensor.to(self.device)) loss.backward(); torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) self.optimizer.step(); return loss.item() @torch.no_grad() def predict_gnn(self, vertices_np, edges, features_np): if not edges: return np.zeros(0, dtype=int) self.model.eval() gnn_input = prepare_gnn_input(vertices_np, edges, features_np) if gnn_input is None or gnn_input['edge_index'].shape[1] == 0: return np.zeros(len(edges), dtype=int) logits = self.model(gnn_input['node_features'].to(self.device), gnn_input['edge_index'].to(self.device), gnn_input['edge_features'].to(self.device)) return logits[:len(edges)].argmax(dim=1).cpu().numpy() @torch.no_grad() def predict(self, features, vertices_np=None, edges=None): if isinstance(self.model, EdgeClassifierGNN) and vertices_np is not None and edges is not None: return self.predict_gnn(vertices_np, edges, features) self.model.eval(); return self.model(torch.from_numpy(features).float().to(self.device)).argmax(dim=1).cpu().numpy() def save(self, path): torch.save({'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.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'])