WireFrameDETR / src /graph_optimizer.py
StarAtNyte1's picture
Add src/graph_optimizer.py
e686396 verified
Raw
History Blame Contribute Delete
2.97 kB
"""Graph topology optimization for S23DR 2026."""
import numpy as np
from typing import List, Tuple
from collections import defaultdict
def resolve_t_junctions(vertices, connections, threshold=0.3):
if len(vertices) < 3 or len(connections) < 1: return vertices, connections
edges_to_remove, edges_to_add = set(), set()
for v_idx in range(len(vertices)):
for e_idx, (a, b) in enumerate(connections):
if v_idx == a or v_idx == b: continue
p, p1, p2 = vertices[v_idx], vertices[a], vertices[b]
edge_vec = p2 - p1; edge_len = np.linalg.norm(edge_vec)
if edge_len < 1e-6: continue
t = np.dot(p - p1, edge_vec / edge_len) / edge_len
if 0.1 < t < 0.9:
proj = p1 + t * edge_vec
if np.linalg.norm(p - proj) < threshold:
edges_to_remove.add((min(a,b), max(a,b)))
edges_to_add.add((min(a, v_idx), max(a, v_idx)))
edges_to_add.add((min(v_idx, b), max(v_idx, b)))
final = [(a, b) for a, b in connections if (min(a,b), max(a,b)) not in edges_to_remove]
existing = set(tuple(sorted(c)) for c in final)
for a, b in edges_to_add:
if (a, b) not in existing: final.append((a, b)); existing.add((a, b))
return vertices, final
def enforce_planarity(vertices, connections, face_vertex_groups=None):
if face_vertex_groups is None: return vertices
adjusted = vertices.copy()
for group in face_vertex_groups:
if len(group) < 3: continue
pts = vertices[group]; centroid = pts.mean(axis=0)
_, _, Vt = np.linalg.svd(pts - centroid); normal = Vt[-1]
for local_idx, global_idx in enumerate(group):
pt = vertices[global_idx]; adjusted[global_idx] = pt - np.dot(pt - centroid, normal) * normal
return adjusted
def merge_close_vertices(vertices, connections, threshold=0.1):
if len(vertices) < 2: return vertices, connections
from scipy.spatial import KDTree
tree = KDTree(vertices); parent = list(range(len(vertices)))
def find(x):
while parent[x] != x: parent[x] = parent[parent[x]]; x = parent[x]
return x
def union(a, b):
ra, rb = find(a), find(b)
if ra != rb: parent[ra] = rb
for i, j in tree.query_pairs(r=threshold): union(i, j)
groups = defaultdict(list)
for i in range(len(vertices)): groups[find(i)].append(i)
new_verts, old_to_new = [], {}
for new_idx, (_, members) in enumerate(groups.items()):
new_verts.append(vertices[members].mean(axis=0))
for m in members: old_to_new[m] = new_idx
new_verts = np.array(new_verts)
new_conns = set()
for a, b in connections:
na, nb = old_to_new.get(a), old_to_new.get(b)
if na is not None and nb is not None and na != nb: new_conns.add(tuple(sorted((na, nb))))
return new_verts, list(new_conns)