| """Deterministic feature encoding of a DrawingGraph for the learned stages. |
| |
| Features are pure functions of vector-exact entity fields (kind, bbox, |
| segment structure, style, text shape) — no pixels, no inference inputs. |
| The classifier additionally sees normalized sheet position; the Siamese |
| matcher does NOT (a moved entity must embed like its counterpart, R7). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
| from cadmorph.models import DrawingEntity, DrawingGraph |
|
|
| KINDS = ("linework", "text", "dimension", "symbol", "hatch") |
|
|
| |
| CLASSES = ("dimension", "door", "furniture", "linework", "note", "room-label", "title", "wall") |
|
|
| BASE_DIM = 17 |
| FEATURE_DIM = BASE_DIM + 2 |
|
|
|
|
| def entity_features( |
| entity: DrawingEntity, graph: DrawingGraph, include_position: bool |
| ) -> list[float]: |
| diag = graph.revision.sheet_diagonal |
| x0, y0, x1, y1 = entity.bbox |
| w, h = x1 - x0, y1 - y0 |
| n_seg = len(entity.geometry) |
| seg_counts = {"line": 0, "rect": 0, "curve": 0, "quad": 0} |
| for seg in entity.geometry: |
| seg_counts[seg.kind] += 1 |
| denom = max(n_seg, 1) |
| text = entity.text_payload or "" |
| style_width = entity.style.width if entity.style and entity.style.width else 0.0 |
|
|
| feats = [ |
| *[1.0 if entity.kind == k else 0.0 for k in KINDS], |
| w / diag, |
| h / diag, |
| min(w / (h + 1e-6), 10.0) / 10.0, |
| min(n_seg, 10) / 10.0, |
| seg_counts["line"] / denom, |
| seg_counts["rect"] / denom, |
| (seg_counts["curve"] + seg_counts["quad"]) / denom, |
| min(style_width, 20.0) / 20.0, |
| 1.0 if text else 0.0, |
| min(len(text), 32) / 32.0, |
| (sum(c.isdigit() for c in text) / len(text)) if text else 0.0, |
| (sum(c.isupper() for c in text) / len(text)) if text else 0.0, |
| ] |
| if include_position: |
| feats += [ |
| ((x0 + x1) / 2.0) / graph.revision.sheet_width, |
| ((y0 + y1) / 2.0) / graph.revision.sheet_height, |
| ] |
| return feats |
|
|
|
|
| def graph_tensors( |
| graph: DrawingGraph, include_position: bool |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """(node features [n, F], edge_index [2, E] undirected) in canonical entity order.""" |
| x = torch.tensor( |
| [entity_features(e, graph, include_position) for e in graph.entities], |
| dtype=torch.float32, |
| ) |
| index = {e.entity_id: i for i, e in enumerate(graph.entities)} |
| pairs: list[tuple[int, int]] = [] |
| for edge in graph.edges: |
| a, b = index[edge.source], index[edge.target] |
| pairs.append((a, b)) |
| pairs.append((b, a)) |
| if pairs: |
| edge_index = torch.tensor(sorted(set(pairs)), dtype=torch.long).t().contiguous() |
| else: |
| edge_index = torch.zeros((2, 0), dtype=torch.long) |
| return x, edge_index |
|
|
