| """Enhanced baseline solution for S23DR 2026.
|
|
|
| Key improvements over the official baseline:
|
| 1. Multi-scale blob detection for vertex finding
|
| 2. Adaptive edge_th based on image resolution
|
| 3. Multi-view edge voting with configurable threshold
|
| 4. COLMAP density-based pruning (not just distance)
|
| 5. Semantic edge classification from Gestalt colors
|
| 6. Better depth fitting with RANSAC instead of median
|
| 7. **Corrupt-image resilience**: per-image try/except for the known
|
| PIL.UnidentifiedImageError near row ~13077 in the training set.
|
| """
|
|
|
| import io
|
| import os
|
| import tempfile
|
| import zipfile
|
| from collections import defaultdict
|
| from typing import Dict, List, Optional, Tuple
|
|
|
| import cv2
|
| import numpy as np
|
| from PIL import Image as PImage, UnidentifiedImageError
|
|
|
| try:
|
| from hoho2025.color_mappings import ade20k_color_mapping, gestalt_color_mapping, EDGE_CLASSES, edge_color_mapping
|
| from hoho2025.example_solutions import (
|
| read_colmap_rec, _cam_matrix_from_image, convert_entry_to_human_readable,
|
| get_house_mask, get_background_mask, get_sparse_depth,
|
| fit_scale_robust_median, get_fitted_dense_depth, get_uv_depth,
|
| project_vertices_to_3d, filter_vertices_by_background,
|
| create_3d_wireframe_single_image,
|
| )
|
| except ImportError:
|
| print("Warning: hoho2025 not installed. Install with: pip install git+https://huggingface.co/usm3d/tools2025.git")
|
| raise
|
|
|
|
|
|
|
|
|
| _VERTEX_MODEL = None
|
| _VERTEX_TRAINER = None
|
| _EDGE_EXIST_TRAINER = None
|
| _EDGE_EXIST_NORM = None
|
|
|
| _BASELINE_DIR = os.path.dirname(os.path.abspath(__file__)) if "__file__" in dir() else "."
|
| _REPO_ROOT = os.path.dirname(_BASELINE_DIR)
|
| _VERTEX_CKPT_PATH = os.path.join(_REPO_ROOT, "checkpoints", "vertex_detector", "best.pt")
|
| _EDGE_EXIST_CKPT_PATH = os.path.join(_REPO_ROOT, "checkpoints", "edge_existence", "best.pt")
|
| _EDGE_EXIST_NORM_PATH = os.path.join(_REPO_ROOT, "checkpoints", "edge_existence", "norm_stats.npz")
|
|
|
|
|
| def _load_vertex_model():
|
| """Lazy-load the trained vertex heatmap model. Returns trainer or None."""
|
| global _VERTEX_MODEL, _VERTEX_TRAINER
|
| if _VERTEX_TRAINER is not None:
|
| return _VERTEX_TRAINER
|
| if not os.path.exists(_VERTEX_CKPT_PATH):
|
| return None
|
| try:
|
| import torch
|
| import sys
|
|
|
| for p in ["/app/src", os.path.join(os.path.dirname(__file__))]:
|
| if p not in sys.path:
|
| sys.path.insert(0, p)
|
| from vertex_detector import VertexHeatmapNet, VertexDetectorTrainer
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| model = VertexHeatmapNet(in_channels=7, num_classes=2, pretrained_backbone=False)
|
| trainer = VertexDetectorTrainer(model, device=device)
|
| trainer.load(_VERTEX_CKPT_PATH)
|
| _VERTEX_TRAINER = trainer
|
| return trainer
|
| except Exception as e:
|
| print(f"[vertex_detector] Could not load model: {e}")
|
| return None
|
|
|
|
|
| def _load_existence_model():
|
| global _EDGE_EXIST_TRAINER, _EDGE_EXIST_NORM
|
| if _EDGE_EXIST_TRAINER is not None:
|
| return _EDGE_EXIST_TRAINER
|
| if not os.path.exists(_EDGE_EXIST_CKPT_PATH):
|
| return None
|
| try:
|
| import torch
|
| import sys
|
| for p in ["/app/src", os.path.dirname(__file__)]:
|
| if p not in sys.path:
|
| sys.path.insert(0, p)
|
| from edge_existence import EdgeExistenceTrainer
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| trainer = EdgeExistenceTrainer(device=device)
|
| trainer.load(_EDGE_EXIST_CKPT_PATH)
|
| _EDGE_EXIST_TRAINER = trainer
|
| if os.path.exists(_EDGE_EXIST_NORM_PATH):
|
| data = np.load(_EDGE_EXIST_NORM_PATH)
|
| _EDGE_EXIST_NORM = {'mean': data['mean'], 'std': data['std']}
|
| print("[edge_existence] Loaded checkpoint")
|
| return trainer
|
| except Exception as e:
|
| print(f"[edge_existence] Could not load model: {e}")
|
| return None
|
|
|
|
|
| def _predict_vertices_learned(gest_img, depth_img, ade_img, threshold=0.3, nms_radius=5):
|
| """Run the learned vertex detector. Returns (vertices_list, vertex_areas_list) or None."""
|
| trainer = _load_vertex_model()
|
| if trainer is None:
|
| return None
|
| try:
|
| from vertex_detector import prepare_input_tensor, extract_vertices_from_heatmap
|
| inp = prepare_input_tensor(gest_img, depth_img, ade_img)
|
| heatmap = trainer.predict(inp)
|
| verts_xy, vtypes = extract_vertices_from_heatmap(heatmap, threshold=threshold, nms_radius=nms_radius)
|
| if len(verts_xy) == 0:
|
| return None
|
| vertices = [{"xy": xy.astype(float), "type": t} for xy, t in zip(verts_xy, vtypes)]
|
| areas = [1.0] * len(vertices)
|
| return vertices, areas
|
| except Exception as e:
|
| print(f"[vertex_detector] Inference error: {e}")
|
| return None
|
|
|
|
|
|
|
|
|
| def _safe_to_numpy(img):
|
| """Convert a PIL Image to uint8 ndarray, returning None on corrupt data."""
|
| try:
|
| arr = np.array(img)
|
| if arr is None or arr.size == 0:
|
| return None
|
| return arr.astype(np.uint8)
|
| except (UnidentifiedImageError, OSError, ValueError):
|
| return None
|
|
|
|
|
| def _point_to_segment_dist(pt, seg_p1, seg_p2):
|
| if np.allclose(seg_p1, seg_p2):
|
| return np.linalg.norm(pt - seg_p1)
|
| seg_vec = seg_p2 - seg_p1
|
| pt_vec = pt - seg_p1
|
| seg_len2 = seg_vec.dot(seg_vec)
|
| t = max(0, min(1, pt_vec.dot(seg_vec) / seg_len2))
|
| proj = seg_p1 + t * seg_vec
|
| return np.linalg.norm(pt - proj)
|
|
|
|
|
|
|
|
|
| def get_vertices_and_edges_enhanced(gest_seg_np, edge_th=15.0,
|
| min_blob_area=4, max_blob_area=5000):
|
| if not isinstance(gest_seg_np, np.ndarray):
|
| gest_seg_np = np.array(gest_seg_np)
|
| vertices, vertex_areas = [], []
|
| for v_class, v_type in [('apex', 'apex'), ('eave_end_point', 'eave_end_point')]:
|
| color = np.array(gestalt_color_mapping[v_class])
|
| mask = cv2.inRange(gest_seg_np, color - 0.5, color + 0.5)
|
| if mask.sum() == 0:
|
| continue
|
| num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
|
| for i in range(1, num_labels):
|
| area = stats[i, cv2.CC_STAT_AREA]
|
| if min_blob_area <= area <= max_blob_area:
|
| vertices.append({"xy": centroids[i], "type": v_type})
|
| vertex_areas.append(area)
|
| if len(vertices) < 2:
|
| return vertices, [], []
|
| apex_pts = np.array([v['xy'] for v in vertices])
|
| connections, edge_types = [], []
|
| for edge_class in ['eave', 'ridge', 'rake', 'valley', 'hip', 'flashing', 'step_flashing']:
|
| if edge_class not in gestalt_color_mapping:
|
| continue
|
| edge_color = np.array(gestalt_color_mapping[edge_class])
|
| mask_raw = cv2.inRange(gest_seg_np, edge_color - 0.5, edge_color + 0.5)
|
| mask = cv2.morphologyEx(mask_raw, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
|
| if mask.sum() == 0:
|
| continue
|
| _, labels, stats, _ = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
|
| for lbl in range(1, labels.max() + 1):
|
| if stats[lbl, cv2.CC_STAT_AREA] < 10:
|
| continue
|
| ys, xs = np.where(labels == lbl)
|
| if len(xs) < 2:
|
| continue
|
| pts_for_fit = np.column_stack([xs, ys]).astype(np.float32)
|
| vx, vy, x0, y0 = cv2.fitLine(pts_for_fit, cv2.DIST_L2, 0, 0.01, 0.01).ravel()
|
| proj = (xs - x0) * vx + (ys - y0) * vy
|
| p1 = np.array([x0 + proj.min() * vx, y0 + proj.min() * vy])
|
| p2 = np.array([x0 + proj.max() * vx, y0 + proj.max() * vy])
|
| dists = np.array([_point_to_segment_dist(apex_pts[j], p1, p2) for j in range(len(apex_pts))])
|
| near = np.where(dists <= edge_th)[0]
|
| if len(near) < 2:
|
| continue
|
| near_pts = apex_pts[near]
|
| i1 = near[np.argmin(np.linalg.norm(near_pts - p1, axis=1))]
|
| i2 = near[np.argmin(np.linalg.norm(near_pts - p2, axis=1))]
|
| if i1 != i2:
|
| ek = tuple(sorted((i1, i2)))
|
| if ek not in [tuple(sorted(c[:2])) for c in connections]:
|
| connections.append((ek[0], ek[1]))
|
| edge_types.append(edge_class)
|
| return vertices, connections, edge_types
|
|
|
|
|
|
|
|
|
| def classify_edges_from_gestalt(vertices_3d, connections, edge_types=None):
|
| if edge_types and len(edge_types) == len(connections):
|
| return [EDGE_CLASSES.get(et, 6) for et in edge_types]
|
| out = []
|
| for a, b in connections:
|
| if a >= len(vertices_3d) or b >= len(vertices_3d):
|
| out.append(6); continue
|
| d = vertices_3d[b] - vertices_3d[a]
|
| l = np.linalg.norm(d)
|
| if l < 1e-6:
|
| out.append(6); continue
|
| vc = abs(d[2] / l)
|
| if vc > 0.7:
|
| out.append(EDGE_CLASSES['rake'])
|
| elif vc < 0.3:
|
| out.append(EDGE_CLASSES['eave'])
|
| else:
|
| out.append(EDGE_CLASSES['hip'])
|
| return out
|
|
|
|
|
|
|
|
|
| def merge_vertices_3d_enhanced(vert_edge_per_image, th=0.5, min_edge_votes=1):
|
| all_3d, conns3d, cur, types = [], [], 0, []
|
| vsrc, esrc, etvotes = {}, defaultdict(set), defaultdict(list)
|
| for ci, data in vert_edge_per_image.items():
|
| verts, conns, v3d = data[0], data[1], data[2]
|
| etypes = data[3] if len(data) == 4 else ['ridge'] * len(conns)
|
| if len(v3d) == 0:
|
| continue
|
| types += [int(v['type'] == 'apex') for v in verts]
|
| all_3d.append(v3d)
|
| for li in range(len(v3d)):
|
| vsrc[cur + li] = ci
|
| for idx, (x, y) in enumerate(conns):
|
| gx, gy = x + cur, y + cur
|
| ek = tuple(sorted((gx, gy)))
|
| esrc[ek].add(ci)
|
| if idx < len(etypes):
|
| etvotes[ek].append(etypes[idx])
|
| conns3d.append((gx, gy))
|
| cur += len(v3d)
|
| if not all_3d:
|
| return np.zeros((2, 3)), [(0, 1)], [6]
|
| all_3d = np.concatenate(all_3d)
|
| if len(all_3d) == 0:
|
| return np.zeros((2, 3)), [(0, 1)], [6]
|
| diff = all_3d[:, None, :] - all_3d[None, :, :]
|
| dm = np.sqrt((diff ** 2).sum(-1))
|
| types = np.array(types)
|
| mm = (dm <= th) & (types[:, None] == types[None, :])
|
| to_merge = sorted(set(tuple(a.nonzero()[0].tolist()) for a in mm))
|
| tmf = defaultdict(list)
|
| for i in range(len(all_3d)):
|
| for j in to_merge:
|
| if i in j:
|
| tmf[i] += j
|
| for k in tmf:
|
| tmf[k] = list(set(tmf[k]))
|
| seen, merged = set(), []
|
| for k, v in tmf.items():
|
| if k not in seen:
|
| merged.append(v)
|
| seen.update(v)
|
| o2n, nv = {}, []
|
| for c, idxs in enumerate(merged):
|
| nv.append(all_3d[idxs].mean(0))
|
| for i in idxs:
|
| o2n[i] = c
|
| nv = np.array(nv)
|
| ne_imgs, ne_tv = defaultdict(set), defaultdict(list)
|
| for conn in conns3d:
|
| na, nb = o2n.get(conn[0]), o2n.get(conn[1])
|
| if na is None or nb is None or na == nb:
|
| continue
|
| nk = tuple(sorted((na, nb)))
|
| ok = tuple(sorted(conn))
|
| ne_imgs[nk] |= esrc.get(ok, set())
|
| ne_tv[nk].extend(etvotes.get(ok, []))
|
| nc, ec = [], []
|
| for ek, imgs in ne_imgs.items():
|
| if len(imgs) >= min_edge_votes:
|
| nc.append(ek)
|
| tv = ne_tv.get(ek, ['ridge'])
|
| from collections import Counter
|
| ec.append(EDGE_CLASSES.get(Counter(tv).most_common(1)[0][0], 6) if tv else 6)
|
| if not nc:
|
| return np.zeros((2, 3)), [(0, 1)], [6]
|
| return nv, nc, ec
|
|
|
|
|
|
|
|
|
| def prune_by_colmap_density(vertices, connections, colmap_rec,
|
| th_dist=3.0, min_nearby_points=3, search_radius=2.0,
|
| edge_classes=None):
|
| """Prune vertices not supported by COLMAP points.
|
|
|
| Returns (vertices, connections, edge_classes) where edge_classes is filtered
|
| to match surviving connections if provided, else returns None.
|
| """
|
| if len(vertices) == 0:
|
| return np.empty((0, 3)), [], edge_classes
|
| try:
|
| sfm = [v.xyz for v in colmap_rec.points3D.values()]
|
| except Exception:
|
| return vertices, connections, edge_classes
|
| if not sfm:
|
| return vertices, connections, edge_classes
|
| sfm = np.array(sfm)
|
| d = np.sqrt(((vertices[:, None, :] - sfm[None, :, :]) ** 2).sum(-1))
|
| mask = (d.min(1) <= th_dist) | ((d <= search_radius).sum(1) >= min_nearby_points)
|
| nv = vertices[mask]
|
| o2n = dict(zip(np.where(mask)[0], range(mask.sum())))
|
| surviving = [(i, a, b) for i, (a, b) in enumerate(connections) if mask[a] and mask[b]]
|
| nc = [(o2n[a], o2n[b]) for _, a, b in surviving]
|
| if edge_classes is not None and len(edge_classes) == len(connections):
|
| ec = [edge_classes[i] for i, _, _ in surviving]
|
| else:
|
| ec = edge_classes
|
| return nv, nc, ec
|
|
|
|
|
| def prune_not_connected(vertices, connections, keep_largest=False, edge_classes=None):
|
| """Remove isolated vertices. Optionally keep only the largest connected component.
|
|
|
| Returns (vertices, connections, edge_classes).
|
| """
|
| if len(vertices) == 0:
|
| return np.array([]), [], edge_classes
|
| used = set()
|
| for i, j in connections:
|
| used.add(i); used.add(j)
|
| if not used:
|
| return np.empty((0, 3)), [], []
|
| if not keep_largest:
|
| ul = sorted(used)
|
| o2n = {o: n for n, o in enumerate(ul)}
|
| surviving = [(i, a, b) for i, (a, b) in enumerate(connections) if a in used and b in used]
|
| nc = [(o2n[a], o2n[b]) for _, a, b in surviving]
|
| ec = ([edge_classes[i] for i, _, _ in surviving]
|
| if edge_classes is not None and len(edge_classes) == len(connections) else edge_classes)
|
| return np.array([vertices[i] for i in ul]), nc, ec
|
| adj = defaultdict(set)
|
| for i, j in connections:
|
| adj[i].add(j); adj[j].add(i)
|
| visited, comps = set(), []
|
| for s in used:
|
| if s in visited:
|
| continue
|
| q, c = [s], []
|
| visited.add(s)
|
| while q:
|
| cur = q.pop()
|
| c.append(cur)
|
| for nb in adj[cur]:
|
| if nb not in visited:
|
| visited.add(nb); q.append(nb)
|
| comps.append(c)
|
| comps.sort(key=len, reverse=True)
|
| lg = set(comps[0]) if comps else set()
|
| o2n = {o: n for n, o in enumerate(sorted(lg))}
|
| nv = np.array([vertices[i] for i in sorted(lg)])
|
| surviving = [(i, a, b) for i, (a, b) in enumerate(connections) if a in lg and b in lg]
|
| nc = list(set(tuple(sorted((o2n[a], o2n[b]))) for _, a, b in surviving))
|
|
|
| ec = classify_edges_from_gestalt(nv, nc) if edge_classes is None else edge_classes
|
| return nv, nc, ec
|
|
|
|
|
|
|
|
|
| def predict_wireframe_enhanced(entry, edge_th=12.0, merge_th=0.5,
|
| prune_dist_th=3.5, min_edge_votes=1,
|
| existence_threshold=0.4, candidate_dist_thresh=8.0,
|
| verbose=False):
|
| """Enhanced wireframe prediction with **corrupt-image resilience**.
|
|
|
| Each per-image processing step is wrapped in try/except so a single
|
| broken image (PIL.UnidentifiedImageError) does not crash the scene.
|
| """
|
| good_entry = convert_entry_to_human_readable(entry)
|
| vert_edge_per_image = {}
|
| colmap_rec = good_entry.get('colmap')
|
|
|
| for i, (gest, depth, img_id, ade_seg) in enumerate(zip(
|
| good_entry['gestalt'], good_entry['depth'],
|
| good_entry['image_ids'], good_entry['ade'],
|
| )):
|
| try:
|
|
|
| depth_np = _safe_to_numpy(depth)
|
| gest_np = _safe_to_numpy(gest)
|
| ade_np = _safe_to_numpy(ade_seg)
|
| if depth_np is None or gest_np is None or ade_np is None:
|
| if verbose:
|
| print(f" ⚠ Skipping image {i} ({img_id}): corrupt/unreadable")
|
| vert_edge_per_image[i] = ([], [], np.empty((0, 3)), [])
|
| continue
|
|
|
| depth_size = (depth_np.shape[1], depth_np.shape[0])
|
| gest_seg = gest.resize(depth_size)
|
| gest_seg_np = np.array(gest_seg).astype(np.uint8)
|
| ade_seg_np = np.array(ade_seg.resize(depth_size)).astype(np.uint8)
|
|
|
|
|
| vertices, connections, edge_types = [], [], []
|
| learned = _predict_vertices_learned(gest, depth, ade_seg, threshold=0.35)
|
| if learned is not None:
|
| heatmap_w, heatmap_h = 128, 96
|
| scale_x = depth_np.shape[1] / heatmap_w
|
| scale_y = depth_np.shape[0] / heatmap_h
|
| for lv in learned[0]:
|
| lxy = np.array([lv['xy'][0] * scale_x, lv['xy'][1] * scale_y])
|
| vertices.append({"xy": lxy, "type": lv["type"]})
|
| if verbose and vertices:
|
| print(f" [learned] view {i}: {len(vertices)} primary vertices")
|
|
|
|
|
| blob_result = get_vertices_and_edges_enhanced(gest_seg_np, edge_th=edge_th)
|
| blob_vertices, blob_connections = blob_result[0], blob_result[1]
|
| blob_edge_types = blob_result[2] if len(blob_result) == 3 else ['ridge'] * len(blob_connections)
|
| existing_pts = [v['xy'] for v in vertices]
|
| for bv in blob_vertices:
|
| if existing_pts and min(
|
| np.linalg.norm(np.array(p) - np.array(bv['xy'])) for p in existing_pts
|
| ) < edge_th:
|
| continue
|
| vertices.append(bv)
|
| existing_pts.append(bv['xy'])
|
|
|
| connections = blob_connections
|
| edge_types = blob_edge_types
|
|
|
| vertices, connections = filter_vertices_by_background(
|
| vertices, connections, ade_seg_np)
|
| edge_types = edge_types[:len(connections)]
|
|
|
| if len(vertices) < 2 or len(connections) < 1:
|
| vert_edge_per_image[i] = ([], [], np.empty((0, 3)), [])
|
| continue
|
|
|
| vertices_3d = 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, edge_types)
|
|
|
| except (UnidentifiedImageError, OSError) as exc:
|
|
|
| if verbose:
|
| print(f" ⚠ Corrupt image in view {i} ({img_id}): {exc}")
|
| vert_edge_per_image[i] = ([], [], np.empty((0, 3)), [])
|
| except Exception as exc:
|
| if verbose:
|
| print(f" ⚠ Error processing view {i} ({img_id}): {exc}")
|
| vert_edge_per_image[i] = ([], [], np.empty((0, 3)), [])
|
|
|
|
|
| all_v, all_c, ecls = merge_vertices_3d_enhanced(
|
| vert_edge_per_image, th=merge_th, min_edge_votes=min_edge_votes)
|
| all_v, all_c, ecls = prune_by_colmap_density(
|
| all_v, all_c, colmap_rec, th_dist=prune_dist_th, edge_classes=ecls)
|
| all_v, all_c, ecls = prune_not_connected(
|
| all_v, all_c, keep_largest=False, edge_classes=ecls)
|
| if len(all_v) < 2:
|
| return np.zeros((2, 3)), [(0, 1)], [6]
|
|
|
|
|
| exist_model = _load_existence_model()
|
| if exist_model is not None and len(all_v) >= 2:
|
| try:
|
| import sys
|
| for p in ["/app/src", os.path.dirname(__file__)]:
|
| if p not in sys.path:
|
| sys.path.insert(0, p)
|
| from edge_existence import generate_candidate_pairs
|
| from edge_classifier import compute_edge_features
|
| candidates = generate_candidate_pairs(all_v, dist_thresh=candidate_dist_thresh)
|
| if candidates:
|
| features = compute_edge_features(all_v, candidates, colmap_rec=colmap_rec)
|
| if _EDGE_EXIST_NORM is not None:
|
| std = _EDGE_EXIST_NORM['std'].copy(); std[std < 1e-6] = 1.0
|
| features = (features - _EDGE_EXIST_NORM['mean']) / std
|
| mask = exist_model.predict(features, threshold=existence_threshold)
|
| all_c = [c for c, m in zip(candidates, mask) if m]
|
| ecls = None
|
| if verbose:
|
| print(f" [existence] {len(candidates)} candidates → {len(all_c)} edges kept")
|
| except Exception as exc:
|
| if verbose:
|
| print(f" [existence] error: {exc}")
|
|
|
| if len(all_c) < 1:
|
| all_c = [(0, 1)]
|
|
|
|
|
| if ecls is None or len(ecls) != len(all_c):
|
| ecls = classify_edges_from_gestalt(all_v, all_c)
|
| else:
|
| geo = classify_edges_from_gestalt(all_v, all_c)
|
| ecls = [e if e != 6 else geo[i] for i, e in enumerate(ecls)]
|
| all_c = [(int(a), int(b)) for a, b in all_c]
|
| return all_v, all_c, ecls
|
|
|
