src/baseline_enhanced.py

"""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