src/edge_classifier.py

"""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'])