Enhance wireframe prediction filtering and clustering

script.py

@@ -11,16 +11,330 @@ import json
 import numpy as np
 from datasets import load_dataset
 from typing import Dict
-from hoho2025.example_solutions import predict_wireframe
 from joblib import Parallel, delayed
+from sklearn.cluster import DBSCAN
+from hoho2025 import example_solutions as hoho_example
+
+VERTEX_MERGE_EPS = 0.5
+EDGE_MIN_LENGTH = 0.5
+EDGE_MIN_SUPPORT_IMAGES = 2
+EDGE_MAX_ANGLE_DEG = 25.0
+VERTEX_MAX_COLMAP_DIST = 4.0
+VERTEX_MIN_EDGE_DEGREE = 2
+VERTEX_MIN_VIEW_COUNT = 2
 
 def empty_solution(sample):
     '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
     return np.zeros((2,3)), [(0, 1)]
 
+def _unit_vector(vector):
+    norm = np.linalg.norm(vector)
+    if norm == 0:
+        return None
+    return vector / norm
+
+def merge_vertices_dbscan(vertices_3d, edges, eps=VERTEX_MERGE_EPS):
+    """Cluster nearby 3D vertices, replace clusters with centroids, and remap edges."""
+    vertices_3d = np.asarray(vertices_3d, dtype=float)
+    if vertices_3d.ndim != 2 or vertices_3d.shape[1] != 3 or len(vertices_3d) == 0:
+        return vertices_3d.reshape((-1, 3)), []
+
+    clustering = DBSCAN(eps=eps, min_samples=1).fit(vertices_3d)
+    labels = clustering.labels_
+    unique_labels = np.unique(labels)
+    label_to_new_idx = {label: idx for idx, label in enumerate(unique_labels)}
+
+    merged_vertices = np.stack(
+        [vertices_3d[labels == label].mean(axis=0) for label in unique_labels],
+        axis=0,
+    )
+
+    remapped_edges = []
+    seen_edges = set()
+    for a, b in edges:
+        a = int(a)
+        b = int(b)
+        if a < 0 or b < 0 or a >= len(labels) or b >= len(labels):
+            continue
+
+        new_a = label_to_new_idx[labels[a]]
+        new_b = label_to_new_idx[labels[b]]
+        if new_a == new_b:
+            continue
+
+        edge = (min(new_a, new_b), max(new_a, new_b))
+        if edge in seen_edges:
+            continue
+
+        seen_edges.add(edge)
+        remapped_edges.append(edge)
+
+    return merged_vertices, remapped_edges
+
+def merge_vertices_dbscan_with_edge_stats(vert_edge_per_image, eps=VERTEX_MERGE_EPS):
+    """Merge multi-view vertices with DBSCAN and keep edge support/direction stats."""
+    all_vertices = []
+    old_edges = []
+    vertex_types = []
+    vertex_src_images = []
+    cur_start = 0
+
+    for img_idx, (vertices, connections, vertices_3d) in vert_edge_per_image.items():
+        vertices_3d = np.asarray(vertices_3d, dtype=float).reshape((-1, 3))
+        if len(vertices_3d) == 0:
+            continue
+
+        all_vertices.append(vertices_3d)
+        vertex_types.extend([int(v.get('type') == 'apex') for v in vertices])
+        vertex_src_images.extend([img_idx] * len(vertices_3d))
+
+        for a, b in connections:
+            a = int(a)
+            b = int(b)
+            if a < 0 or b < 0 or a >= len(vertices_3d) or b >= len(vertices_3d):
+                continue
+            old_edges.append((cur_start + a, cur_start + b, img_idx))
+
+        cur_start += len(vertices_3d)
+
+    if not all_vertices:
+        return np.empty((0, 3)), [], np.array([], dtype=int), {}, {}
+
+    all_vertices = np.concatenate(all_vertices, axis=0)
+    labels = DBSCAN(eps=eps, min_samples=1).fit(all_vertices).labels_
+
+    # Keep apex and non-apex vertices from collapsing into the same final corner.
+    vertex_types = np.asarray(vertex_types, dtype=int)
+    cluster_keys = [(int(label), int(vtype)) for label, vtype in zip(labels, vertex_types)]
+    unique_keys = sorted(set(cluster_keys))
+    key_to_new_idx = {key: idx for idx, key in enumerate(unique_keys)}
+    old_to_new_idx = np.array([key_to_new_idx[key] for key in cluster_keys], dtype=int)
+
+    merged_vertices = np.stack(
+        [
+            all_vertices[old_to_new_idx == idx].mean(axis=0)
+            for idx in range(len(unique_keys))
+        ],
+        axis=0,
+    )
+
+    vertex_src_images = np.asarray(vertex_src_images)
+    vertex_view_count = np.array(
+        [
+            len(set(vertex_src_images[old_to_new_idx == idx]))
+            for idx in range(len(unique_keys))
+        ],
+        dtype=int,
+    )
+
+    edge_image_sets = {}
+    edge_dirs = {}
+    for old_a, old_b, img_idx in old_edges:
+        new_a = int(old_to_new_idx[old_a])
+        new_b = int(old_to_new_idx[old_b])
+        if new_a == new_b:
+            continue
+
+        edge = (min(new_a, new_b), max(new_a, new_b))
+        direction = _unit_vector(all_vertices[old_b] - all_vertices[old_a])
+        if direction is None:
+            continue
+
+        final_direction = _unit_vector(merged_vertices[edge[1]] - merged_vertices[edge[0]])
+        if final_direction is None:
+            continue
+        if np.dot(direction, final_direction) < 0:
+            direction = -direction
+
+        edge_image_sets.setdefault(edge, set()).add(img_idx)
+        edge_dirs.setdefault(edge, []).append(direction)
+
+    edge_vote_count = {edge: len(imgs) for edge, imgs in edge_image_sets.items()}
+    edge_angle_ok = {
+        edge: edge_directions_are_consistent(dirs, EDGE_MAX_ANGLE_DEG)
+        for edge, dirs in edge_dirs.items()
+    }
+
+    return merged_vertices, list(edge_image_sets.keys()), vertex_view_count, edge_vote_count, edge_angle_ok
+
+def edge_directions_are_consistent(directions, max_angle_deg=EDGE_MAX_ANGLE_DEG):
+    """Check whether per-image 3D edge directions agree after sign alignment."""
+    if len(directions) <= 1:
+        return True
+
+    directions = np.asarray(directions, dtype=float)
+    mean_direction = _unit_vector(directions.mean(axis=0))
+    if mean_direction is None:
+        return False
+
+    min_cos = np.cos(np.deg2rad(max_angle_deg))
+    return bool(np.all(directions @ mean_direction >= min_cos))
+
+def filter_edges_by_geometry(vertices_3d, edges, edge_vote_count=None, edge_angle_ok=None,
+                             min_len=EDGE_MIN_LENGTH, min_support=EDGE_MIN_SUPPORT_IMAGES):
+    """Drop short, weakly supported, or direction-inconsistent edges."""
+    vertices_3d = np.asarray(vertices_3d, dtype=float)
+    if len(vertices_3d) == 0:
+        return []
+
+    filtered_edges = []
+    seen_edges = set()
+    for a, b in edges:
+        a = int(a)
+        b = int(b)
+        if a < 0 or b < 0 or a >= len(vertices_3d) or b >= len(vertices_3d) or a == b:
+            continue
+
+        edge = (min(a, b), max(a, b))
+        if edge in seen_edges:
+            continue
+
+        edge_length = np.linalg.norm(vertices_3d[edge[1]] - vertices_3d[edge[0]])
+        if edge_length < min_len:
+            continue
+
+        if edge_vote_count is not None and edge_vote_count.get(edge, 0) < min_support:
+            continue
+
+        if edge_angle_ok is not None and not edge_angle_ok.get(edge, True):
+            continue
+
+        seen_edges.add(edge)
+        filtered_edges.append(edge)
+
+    return filtered_edges
+
+def prune_bad_vertices(vertices_3d, edges, colmap_rec, vertex_view_count,
+                       max_colmap_dist=VERTEX_MAX_COLMAP_DIST,
+                       min_edge_degree=VERTEX_MIN_EDGE_DEGREE,
+                       min_view_count=VERTEX_MIN_VIEW_COUNT):
+    """Remove vertices with weak COLMAP, edge-degree, or multi-view support."""
+    vertices_3d = np.asarray(vertices_3d, dtype=float)
+    vertex_view_count = np.asarray(vertex_view_count, dtype=int)
+    if len(vertices_3d) == 0:
+        return vertices_3d.reshape((-1, 3)), []
+
+    valid_edges = []
+    for a, b in edges:
+        a = int(a)
+        b = int(b)
+        if 0 <= a < len(vertices_3d) and 0 <= b < len(vertices_3d) and a != b:
+            valid_edges.append((a, b))
+
+    colmap_mask = np.ones(len(vertices_3d), dtype=bool)
+    xyz_sfm = [point.xyz for point in colmap_rec.points3D.values()]
+    if xyz_sfm:
+        xyz_sfm = np.asarray(xyz_sfm, dtype=float)
+        diff = vertices_3d[:, None, :] - xyz_sfm[None, :, :]
+        min_dist = np.sqrt((diff ** 2).sum(axis=-1)).min(axis=1)
+        colmap_mask = min_dist <= max_colmap_dist
+
+    view_mask = vertex_view_count >= min_view_count
+    keep_mask = colmap_mask & view_mask
+
+    while True:
+        degree = np.zeros(len(vertices_3d), dtype=int)
+        for a, b in valid_edges:
+            if keep_mask[a] and keep_mask[b]:
+                degree[a] += 1
+                degree[b] += 1
+
+        next_keep_mask = keep_mask & (degree >= min_edge_degree)
+        if np.array_equal(next_keep_mask, keep_mask):
+            break
+        keep_mask = next_keep_mask
+
+    old_to_new = {}
+    new_vertices = []
+    for old_idx, keep in enumerate(keep_mask):
+        if keep:
+            old_to_new[old_idx] = len(new_vertices)
+            new_vertices.append(vertices_3d[old_idx])
+
+    new_edges = []
+    seen_edges = set()
+    for a, b in valid_edges:
+        if a not in old_to_new or b not in old_to_new:
+            continue
+        edge = tuple(sorted((old_to_new[a], old_to_new[b])))
+        if edge in seen_edges:
+            continue
+        seen_edges.add(edge)
+        new_edges.append(edge)
+
+    if not new_vertices:
+        return np.empty((0, 3)), []
+    return np.asarray(new_vertices), new_edges
+
+def predict_wireframe_filtered(sample, verbose=False):
+    """Baseline prediction with DBSCAN vertex merge plus stricter edge filtering."""
+    good_entry = hoho_example.convert_entry_to_human_readable(sample)
+    vert_edge_per_image = {}
+    colmap_rec = None
+
+    for i, (gest, depth, img_id, ade_seg) in enumerate(zip(
+        good_entry['gestalt'],
+        good_entry['depth'],
+        good_entry['image_ids'],
+        good_entry['ade'],
+    )):
+        if 'colmap' in good_entry:
+            colmap_rec = good_entry['colmap']
+        else:
+            colmap_rec = good_entry['colmap_binary']
+
+        depth_size = (np.array(depth).shape[1], np.array(depth).shape[0])
+        gest_seg_np = np.array(gest.resize(depth_size)).astype(np.uint8)
+        vertices, connections = hoho_example.get_vertices_and_edges_from_segmentation(
+            gest_seg_np, edge_th=10.
+        )
+
+        ade_seg_np = np.array(ade_seg.resize(depth_size)).astype(np.uint8)
+        vertices, connections = hoho_example.filter_vertices_by_background(
+            vertices, connections, ade_seg_np
+        )
+
+        if (len(vertices) < 2) or (len(connections) < 1):
+            if verbose:
+                print(f'Not enough vertices or connections found in image {i}, skipping.')
+            vert_edge_per_image[i] = [], [], np.empty((0, 3))
+            continue
+
+        vertices_3d = hoho_example.create_3d_wireframe_single_image(
+            vertices, connections, depth, colmap_rec, img_id, ade_seg, verbose=verbose
+        )
+        vert_edge_per_image[i] = vertices, connections, vertices_3d
+
+    all_3d_vertices, connections_3d, vertex_view_count, edge_vote_count, edge_angle_ok = (
+        merge_vertices_dbscan_with_edge_stats(vert_edge_per_image, VERTEX_MERGE_EPS)
+    )
+    connections_3d = filter_edges_by_geometry(
+        all_3d_vertices,
+        connections_3d,
+        edge_vote_count=edge_vote_count,
+        edge_angle_ok=edge_angle_ok,
+    )
+    if colmap_rec is None or len(all_3d_vertices) < 2 or len(connections_3d) < 1:
+        return empty_solution(sample)
+
+    all_3d_vertices_clean, connections_3d_clean = prune_bad_vertices(
+        all_3d_vertices, connections_3d, colmap_rec, vertex_view_count
+    )
+    all_3d_vertices_clean, connections_3d_clean = hoho_example.prune_not_connected(
+        all_3d_vertices_clean, connections_3d_clean, keep_largest=False
+    )
+
+    if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1:
+        if verbose:
+            print(f'Not enough vertices or connections in the 3D vertices')
+        return empty_solution(sample)
+
+    connections_3d_clean = [(int(a), int(b)) for a, b in connections_3d_clean]
+    return all_3d_vertices_clean, connections_3d_clean
+
 def predict_wireframe_safely(sample):
     try:
-        pred_vertices, pred_edges = predict_wireframe(sample)
+        pred_vertices, pred_edges = predict_wireframe_filtered(sample)
     except Exception as e:
         print (f"Failed due to {e},  returning empty solution")
         pred_vertices, pred_edges = empty_solution(sample)
@@ -155,4 +469,4 @@ if __name__ == "__main__":
     with open("submission.json", "w") as f:
         json.dump(solution, f)
 
-    print("------------ Done ------------ ")
+    print("------------ Done ------------ ")