File size: 11,478 Bytes
76b06d9 99819e3 76b06d9 99819e3 76b06d9 99819e3 76b06d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | """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'])
|