| """V4 vertex classifier with DINOv2 patch features. |
| |
| Per-vertex analog of edge_classifier_v4. For each predicted 3D vertex, |
| project to each view, bilinearly sample DINOv2 patch features at the |
| projected pixel location. Mean+max pool across views -> 768-dim. Concat |
| with simple geometric features (10-dim) -> 778-dim. MLP head -> P(keep). |
| |
| Decisions: |
| - drop a vertex if classifier P(keep) < threshold |
| - re-run drop_orphan after to clean dangling edges |
| |
| Label for training: vertex is "true" if there's a GT vertex within |
| `match_radius` meters. |
| |
| Vertex geometric features (10): |
| 0 z (height) |
| 1 degree (number of incident edges) |
| 2 dist to nearest other vertex |
| 3 dist to nearest colmap point |
| 4 num views vertex is in-frame |
| 5 min projected dist to nearest gestalt corner pixel (clipped 30) |
| 6 mean projected dist to nearest gestalt corner pixel |
| 7 num views with corner pixel within 5px |
| 8 num views with corner pixel within 15px |
| 9 vertex_count / median scene vertex count (graph density hint) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| from edge_classifier_v4 import ( |
| DINO_FEAT_DIM, get_dino_model, _encode_views_with_dino, _bilinear_sample_grid, |
| ) |
|
|
| POINT_CLASSES = ("apex", "eave_end_point", "flashing_end_point") |
| V_GEOM_DIM = 10 |
| V_EDGE_FEAT_DIM = DINO_FEAT_DIM * 2 |
|
|
|
|
| class VertexClassifierV4(nn.Module): |
| def __init__(self, geom_dim: int = V_GEOM_DIM, edge_feat_dim: int = V_EDGE_FEAT_DIM, |
| hidden: int = 128): |
| super().__init__() |
| in_dim = geom_dim + edge_feat_dim |
| self.net = nn.Sequential( |
| nn.Linear(in_dim, hidden), |
| nn.GELU(), |
| nn.Dropout(0.2), |
| nn.Linear(hidden, hidden), |
| nn.GELU(), |
| nn.Dropout(0.2), |
| nn.Linear(hidden, hidden // 2), |
| nn.GELU(), |
| nn.Linear(hidden // 2, 1), |
| ) |
|
|
| def forward(self, geom_feats, dino_feats): |
| x = torch.cat([geom_feats, dino_feats], dim=1) |
| return self.net(x).squeeze(-1) |
|
|
|
|
| def _build_corner_dt(good, views, dilate_px=0): |
| """Per-view corner DT (distance to nearest gestalt point-class pixel).""" |
| import cv2 |
| from hoho2025.color_mappings import gestalt_color_mapping |
| out = {} |
| for gest_pil, depth_pil, img_id in zip( |
| good["gestalt"], good["depth"], good["image_ids"] |
| ): |
| if img_id not in views: |
| continue |
| depth_np = np.array(depth_pil) |
| H, W = depth_np.shape[:2] |
| gest_np = np.array(gest_pil.resize((W, H))).astype(np.uint8) |
| mask = np.zeros((H, W), dtype=np.uint8) |
| for cls in POINT_CLASSES: |
| color = np.array(gestalt_color_mapping[cls]) |
| mask |= cv2.inRange(gest_np, color - 0.5, color + 0.5) |
| if mask.sum() > 0: |
| dt = cv2.distanceTransform(255 - mask, cv2.DIST_L2, 5) |
| dt = np.minimum(dt, 30.0).astype(np.float32) |
| else: |
| dt = np.full((H, W), 30.0, dtype=np.float32) |
| out[img_id] = (dt, H, W) |
| return out |
|
|
|
|
| def _vertex_geom_features(pv, pe, sample): |
| """Per-vertex geometric features (V, 10).""" |
| pv_arr = np.asarray(pv, dtype=np.float32) |
| V = len(pv_arr) |
| feats = np.zeros((V, V_GEOM_DIM), dtype=np.float32) |
| if V == 0: |
| return feats |
|
|
| deg = np.zeros(V, dtype=np.int32) |
| for a, b in pe: |
| if 0 <= int(a) < V: deg[int(a)] += 1 |
| if 0 <= int(b) < V: deg[int(b)] += 1 |
|
|
| try: |
| from hoho2025.example_solutions import convert_entry_to_human_readable |
| from mvs_utils import collect_views, project_world_to_image |
| from scipy.spatial import cKDTree |
|
|
| good = convert_entry_to_human_readable(sample) |
| colmap_rec = good.get("colmap") or good.get("colmap_binary") |
| if colmap_rec is None: |
| return feats |
| views = collect_views(colmap_rec, good["image_ids"]) |
| if not views: |
| return feats |
| corner_dt = _build_corner_dt(good, views) |
|
|
| |
| c_pts = [] |
| if hasattr(colmap_rec, "points3D"): |
| for p in colmap_rec.points3D.values(): |
| c_pts.append(p.xyz) |
| c_tree = cKDTree(np.asarray(c_pts, dtype=np.float32)) if c_pts else None |
|
|
| |
| pv_tree = cKDTree(pv_arr) if V >= 2 else None |
|
|
| for i in range(V): |
| v = pv_arr[i] |
| feats[i, 0] = float(v[2]) |
| feats[i, 1] = float(deg[i]) |
| if pv_tree is not None: |
| dists, _ = pv_tree.query(v, k=min(2, V)) |
| feats[i, 2] = float(dists[1] if len(dists) > 1 else 0.0) |
| if c_tree is not None: |
| d, _ = c_tree.query(v) |
| feats[i, 3] = float(d) |
|
|
| |
| view_dists = [] |
| in_frame = 0 |
| close_5 = 0 |
| close_15 = 0 |
| for img_id, view in views.items(): |
| if img_id not in corner_dt: |
| continue |
| dt, H, W = corner_dt[img_id] |
| uv, z = project_world_to_image(view["P"], v.reshape(1, 3)) |
| if z[0] <= 0: |
| continue |
| if not (0 <= uv[0, 0] < W and 0 <= uv[0, 1] < H): |
| continue |
| in_frame += 1 |
| cd = float(dt[int(uv[0, 1]), int(uv[0, 0])]) |
| view_dists.append(cd) |
| if cd < 5: close_5 += 1 |
| if cd < 15: close_15 += 1 |
| feats[i, 4] = float(in_frame) |
| if view_dists: |
| arr = np.asarray(view_dists) |
| feats[i, 5] = float(arr.min()) |
| feats[i, 6] = float(arr.mean()) |
| else: |
| feats[i, 5] = 30.0 |
| feats[i, 6] = 30.0 |
| feats[i, 7] = float(close_5) |
| feats[i, 8] = float(close_15) |
| feats[i, 9] = float(V) |
| except Exception: |
| pass |
| return feats |
|
|
|
|
| def extract_vertex_features_v4(pv, pe, sample, dino, device="cpu"): |
| """Return (V, 10) geom + (V, 768) dino features per vertex.""" |
| pv_arr = np.asarray(pv) |
| V = len(pv_arr) |
| geom = _vertex_geom_features(pv, pe, sample) |
| dino_feats = np.zeros((V, V_EDGE_FEAT_DIM), dtype=np.float32) |
| if V == 0: |
| return geom, dino_feats |
|
|
| try: |
| from hoho2025.example_solutions import convert_entry_to_human_readable |
| from mvs_utils import collect_views, project_world_to_image |
| good = convert_entry_to_human_readable(sample) |
| colmap_rec = good.get("colmap") or good.get("colmap_binary") |
| if colmap_rec is None: |
| return geom, dino_feats |
| views = collect_views(colmap_rec, good["image_ids"]) |
| if not views: |
| return geom, dino_feats |
| per_view, Hs, Ws = _encode_views_with_dino(good, views, dino, device) |
| if not per_view: |
| return geom, dino_feats |
|
|
| for i in range(V): |
| v = pv_arr[i].reshape(1, 3) |
| per_view_feats = [] |
| for img_id, view in views.items(): |
| if img_id not in per_view: |
| continue |
| H, W = Hs[img_id], Ws[img_id] |
| uv, z = project_world_to_image(view["P"], v) |
| if z[0] <= 0: |
| continue |
| u, vu = uv[0, 0], uv[0, 1] |
| if not (0 <= u < W and 0 <= vu < H): |
| continue |
| feat = _bilinear_sample_grid(per_view[img_id], u / max(W - 1, 1), vu / max(H - 1, 1)) |
| per_view_feats.append(feat) |
| if per_view_feats: |
| arr = np.asarray(per_view_feats) |
| dino_feats[i] = np.concatenate([arr.mean(axis=0), arr.max(axis=0)]) |
| except Exception: |
| pass |
| return geom, dino_feats |
|
|
|
|
| def label_vertices_vs_gt(pv, gt_v, match_radius: float = 0.5): |
| """Vertex is positive if a GT vertex is within match_radius meters.""" |
| pv_arr = np.asarray(pv, dtype=np.float32) |
| gt_v_arr = np.asarray(gt_v, dtype=np.float32) |
| if pv_arr.shape[0] == 0 or gt_v_arr.shape[0] == 0: |
| return np.zeros(pv_arr.shape[0], dtype=np.float32) |
| from scipy.spatial import cKDTree |
| tree = cKDTree(gt_v_arr) |
| dists, _ = tree.query(pv_arr) |
| return (dists <= match_radius).astype(np.float32) |
|
|
|
|
| def classify_vertices_v4(pv, pe, sample, classifier, dino, device="cpu", |
| threshold: float = 0.5, |
| feature_mean=None, feature_std=None, |
| edge_feat_mean=None, edge_feat_std=None, |
| min_keep_frac: float = 0.85): |
| """Drop low-confidence vertices, rebuild edge indexing.""" |
| try: |
| pv_arr = np.asarray(pv) |
| if len(pv_arr) == 0: |
| return pv, pe |
| geom, dino_feats = extract_vertex_features_v4(pv, pe, sample, dino, device=device) |
| if feature_mean is not None and feature_std is not None: |
| geom = (geom - feature_mean) / (feature_std + 1e-6) |
| if edge_feat_mean is not None and edge_feat_std is not None: |
| dino_feats = (dino_feats - edge_feat_mean) / (edge_feat_std + 1e-6) |
| with torch.no_grad(): |
| g = torch.tensor(geom, dtype=torch.float32) |
| d = torch.tensor(dino_feats, dtype=torch.float32) |
| scores = torch.sigmoid(classifier(g, d)).numpy() |
|
|
| keep_mask = scores >= threshold |
| min_keep = max(2, int(np.ceil(min_keep_frac * len(pv_arr)))) |
| if keep_mask.sum() < min_keep: |
| top_idx = np.argsort(-scores)[:min_keep] |
| keep_mask = np.zeros_like(keep_mask) |
| keep_mask[top_idx] = True |
|
|
| if keep_mask.all(): |
| return pv, pe |
|
|
| keep_idx = np.where(keep_mask)[0] |
| old_to_new = {int(o): int(n) for n, o in enumerate(keep_idx)} |
| new_pv = pv_arr[keep_idx] |
| new_pe = [(old_to_new[int(a)], old_to_new[int(b)]) |
| for a, b in pe |
| if int(a) in old_to_new and int(b) in old_to_new] |
| if len(new_pv) < 2 or len(new_pe) < 1: |
| return pv, pe |
| return new_pv, new_pe |
| except Exception: |
| return pv, pe |
|
|
|
|
| def load_classifier_v4(path: str, device: str = "cpu"): |
| blob = torch.load(path, map_location=device, weights_only=False) |
| m = VertexClassifierV4(hidden=blob.get("hidden", 128)) |
| m.load_state_dict(blob["model"]) |
| m.to(device).eval() |
| return (m, blob.get("feature_mean"), blob.get("feature_std"), |
| blob.get("edge_feat_mean"), blob.get("edge_feat_std")) |
|
|
