"""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 # ── lazy-loaded models ──────────────────────────────────────────────────── _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 # Support both /app/src and local src 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 # ── helpers ──────────────────────────────────────────────────────────────── 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) # ── 2-D vertex / edge extraction ────────────────────────────────────────── 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 # ── semantic classification ──────────────────────────────────────────────── 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 # ── multi-view merge ────────────────────────────────────────────────────── 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 # ── pruning ──────────────────────────────────────────────────────────────── 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)) # edge_classes tracking through set-dedup is lossy; fall back to geometric for keep_largest ec = classify_edges_from_gestalt(nv, nc) if edge_classes is None else edge_classes return nv, nc, ec # ── main pipeline ───────────────────────────────────────────────────────── 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: # ── safety: force-decode every image ────────────────────── 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]) # (W, H) 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) # Primary: learned heatmap (lower threshold for better recall) 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") # Secondary: gestalt blobs — only if no heatmap vertex nearby 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']) # Use blob edges as initial 2D connectivity (existence classifier refines later) 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: # Known corrupt-image error — skip this view silently 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)), []) # ── merge + prune ───────────────────────────────────────────────── 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] # ── edge existence: replace gestalt edges with learned cross-view edges ── 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 # re-classify below 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)] # ── semantic classification ─────────────────────────────────────── 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