Upload s23dr_2026/inference.py with huggingface_hub

s23dr_2026/inference.py

@@ -0,0 +1,262 @@
+import math
+
+import numpy as np
+import torch
+from hoho2025.color_mappings import gestalt_color_mapping
+
+from .dataset import sample_points, colors_to_features, combined_sampling, compute_norm_params
+from .model import WireframeBase
+from .scene import Scene
+from .postprocess import (
+    add_edges_by_translated_edge_symmetry,
+    add_gap_fill_edges_from_candidates,
+    filter_edges_by_length,
+    merge_vertices_with_features,
+    post_process_wireframe,
+    project_near_horizontal_edges_to_xz_plane,
+    remove_solitary_edges as remove_solitary_edges_from_wireframe,
+    snap_edges_to_principal_xz_directions,
+)
+
+
+def _rotation_matrix_y(theta: float) -> np.ndarray:
+    c, s = np.cos(theta), np.sin(theta)
+    return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float32)
+
+
+def get_colmap_points(scene: Scene) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    assert len(scene.data["ade"]) == len(scene.data["gestalt"]) == scene.num_images
+    ade = [scene.ade(idx) for idx in range(scene.num_images)]
+    gestalt = [scene.gestalt(idx) for idx in range(scene.num_images)]
+
+    reconstruction = scene.reconstruction
+    xyz = np.stack([pt3d.xyz for pt3d in reconstruction.points3D.values()])
+    rgb = np.stack([pt3d.color for pt3d in reconstruction.points3D.values()])
+
+    R, t, K = [], [], []
+    image_id_to_idx = {}
+    for idx in range(scene.num_images):
+        image = reconstruction.find_image_with_name(scene.image_name(idx))
+        cam_from_world = image.cam_from_world
+        if callable(cam_from_world):
+            cam_from_world = cam_from_world()
+        R.append(cam_from_world.rotation.matrix())
+        t.append(cam_from_world.translation)
+        camera = reconstruction.cameras[image.camera_id]
+        K.append(camera.calibration_matrix())
+        image_id_to_idx[image.image_id] = idx
+
+        ade_width, ade_height = ade[idx].shape[1], ade[idx].shape[0]
+        gestalt_width, gestalt_height = gestalt[idx].shape[1], gestalt[idx].shape[0]
+        if ade_width != gestalt_width or ade_height != gestalt_height:
+            raise ValueError(f"ADE/Gestalt dimension mismatch for image {image.name} in scene {scene.order_id}")
+        if abs(ade_width - camera.width) > 1 or abs(ade_height - camera.height) > 1:
+            raise ValueError(f"Camera/ADE dimension mismatch for image {image.name} in scene {scene.order_id}")
+    R, t, K = np.stack(R), np.stack(t), np.stack(K)
+
+    X = (xyz @ R.transpose(0, 2, 1) + t[:, None]) @ K.transpose(0, 2, 1)
+    xy = X[..., :2] / X[..., 2:3]
+    xy = np.round(xy).astype(int)
+    ade_size = np.array([ade[idx].shape for idx in range(scene.num_images)])
+    xy = np.clip(xy, 0, ade_size[:, None, [1, 0]] - 1)
+
+    def most_common_entry(rgbs):
+        counts = {}
+        for rgb_val in rgbs:
+            key = tuple(rgb_val)
+            counts[key] = counts.get(key, 0) + 1
+        return np.array(max(counts, key=counts.get))
+
+    ade_rgb, gestalt_rgb = [], []
+    for point_idx, pt3d in enumerate(reconstruction.points3D.values()):
+        ade_rgbs, gestalt_rgbs = [], []
+        for elem in pt3d.track.elements:
+            idx = image_id_to_idx[elem.image_id]
+            x, y = xy[idx, point_idx]
+            ade_rgbs.append(ade[idx][y, x])
+            gestalt_rgbs.append(gestalt[idx][y, x])
+        ade_rgb.append(most_common_entry(ade_rgbs))
+        gestalt_rgb.append(most_common_entry(gestalt_rgbs))
+    ade_rgb = np.stack(ade_rgb)
+    gestalt_rgb = np.stack(gestalt_rgb)
+
+    mask_colors = np.array([gestalt_color_mapping[cls] for cls in ("unclassified", "unknown")], dtype=np.uint8)
+    mask = np.any(np.all(gestalt_rgb[:, None, :] == mask_colors[None, :, :], axis=-1), axis=1)
+    return xyz[~mask], rgb[~mask], ade_rgb[~mask], gestalt_rgb[~mask]
+
+
+def predict_wireframe_v2(
+    scene: Scene,
+    model: WireframeBase,
+    device: str,
+    pt_type: str = "colmap",
+    num_points: int = 7168,
+    score_threshold: float = 0.9,
+    merge_distance_threshold: float = 0.5,
+    min_edge_length: float = 0.05,
+    max_edge_length: float = 100.0,
+    remove_solitary_edges: bool = False,
+    filter_by_length: bool = False,
+    project_near_horizontal_edges: bool = False,
+    snap_to_principal_directions: bool = False,
+    horizontal_angle_threshold_deg: float = 5.0,
+    principal_direction_angle_threshold_deg: float = 5.0,
+
+    intersection_angle_threshold_deg: float | None = None,
+    gap_fill_edges: bool = False,
+    gap_fill_score_threshold: float = 0.2,
+    complete_symmetric_edges: bool = False,
+    symmetric_edge_distance_threshold: float = 0.2,
+    merge_method: str = "distance",
+    feature_alpha: float = 0.95,
+) -> tuple[np.ndarray, np.ndarray]:
+    assert pt_type in ("colmap", "monodepth", "both")
+
+    if pt_type == "colmap":
+        points, colors, ade_colors, gestalt_colors = get_colmap_points(scene)
+        xyz = torch.from_numpy(points).float()
+        features = colors_to_features(gestalt_colors, ade_colors, rgb=colors, one_hot=False)
+        min_xyz, max_xyz, q_diff = compute_norm_params(xyz)
+        xyz = (xyz - min_xyz) / q_diff
+        indices = torch.from_numpy(sample_points(gestalt_colors, num_points, weight_vertex=100.0, weight_edge=10.0))
+        xyz = xyz[indices]
+        features = features[indices]
+    else:
+        raise ValueError("Only pt_type='colmap' supported in this implementation.")
+
+    xyz, features = xyz.to(device), features.to(device)
+    do_gap_fill = gap_fill_edges and gap_fill_score_threshold < score_threshold
+    inference_threshold = gap_fill_score_threshold if do_gap_fill else score_threshold
+    results = model.inference(xyz, features, threshold=inference_threshold)
+
+    verts, edges = results["vertices"], results["edges"]
+    min_xyz_np, max_xyz_np = min_xyz.numpy(), max_xyz.numpy()
+    verts = verts * (max_xyz_np - min_xyz_np) + min_xyz_np
+
+    candidate_edge_points = None
+    if do_gap_fill:
+        scores = results["scores"]
+        edge_points = verts[edges]
+        high_confidence = scores > score_threshold
+        gap_fill_candidates = (scores > gap_fill_score_threshold) & (scores <= score_threshold)
+        verts = edge_points[high_confidence].reshape(-1, 3)
+        edges = np.arange(len(verts), dtype=edges.dtype).reshape(-1, 2)
+        candidate_edge_points = edge_points[gap_fill_candidates]
+
+    if merge_method == "auto":
+        merge_method = "feature" if "vertex_feats" in results else "distance"
+    if merge_method == "feature" and "vertex_feats" not in results:
+        merge_method = "distance"
+
+    if merge_method == "feature":
+        vertex_feats = np.asarray(results["vertex_feats"], dtype=np.float32)
+        edge_scores = np.asarray(results["scores"], dtype=np.float32)
+        if do_gap_fill:
+            vertex_feats = vertex_feats.reshape(-1, 2, vertex_feats.shape[-1])[high_confidence].reshape(-1, vertex_feats.shape[-1])
+            edge_scores = edge_scores[high_confidence]
+        verts, edges, _ = merge_vertices_with_features(
+            verts, edges, vertex_feats,
+            alpha=feature_alpha,
+            distance_threshold=merge_distance_threshold,
+            vertex_scores=np.repeat(edge_scores, 2),
+        )
+        if complete_symmetric_edges:
+            verts, edges = add_edges_by_translated_edge_symmetry(verts, edges, distance_threshold=symmetric_edge_distance_threshold)
+        if filter_by_length:
+            verts, edges = filter_edges_by_length(verts, edges, min_edge_length, max_edge_length)
+        if project_near_horizontal_edges:
+            verts, edges = project_near_horizontal_edges_to_xz_plane(verts, edges, horizontal_angle_threshold_deg=horizontal_angle_threshold_deg, principal_direction_angle_threshold_deg=principal_direction_angle_threshold_deg)
+        if snap_to_principal_directions:
+            verts, edges = snap_edges_to_principal_xz_directions(verts, edges, horizontal_angle_threshold_deg=horizontal_angle_threshold_deg, principal_direction_angle_threshold_deg=principal_direction_angle_threshold_deg)
+        if remove_solitary_edges:
+            verts, edges = remove_solitary_edges_from_wireframe(verts, edges)
+        if candidate_edge_points is not None:
+            verts, edges = add_gap_fill_edges_from_candidates(verts, edges, candidate_edge_points, merge_distance_threshold)
+    else:
+        verts, edges = post_process_wireframe(
+            verts, edges, merge_distance_threshold,
+            intersection_angle_threshold_deg=intersection_angle_threshold_deg,
+            remove_solitary_edges=remove_solitary_edges,
+            min_edge_length=min_edge_length,
+            max_edge_length=max_edge_length,
+            filter_by_length=filter_by_length,
+            project_near_horizontal_edges=project_near_horizontal_edges,
+            snap_to_principal_directions=snap_to_principal_directions,
+            horizontal_angle_threshold_deg=horizontal_angle_threshold_deg,
+            principal_direction_angle_threshold_deg=principal_direction_angle_threshold_deg,
+            gap_fill_candidate_edge_points=candidate_edge_points,
+            complete_symmetric_edges=complete_symmetric_edges,
+            symmetric_edge_distance_threshold=symmetric_edge_distance_threshold,
+        )
+
+    return verts, edges
+
+
+def predict_wireframe_v2_tta(
+    scene: Scene,
+    model: WireframeBase,
+    device: str,
+    pt_type: str = "colmap",
+    num_points: int = 7168,
+    score_threshold: float = 0.9,
+    merge_distance_threshold: float = 0.5,
+
+    n_rotations: int = 4,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Test-Time Augmentation: predict over N evenly-spaced Y-axis rotations and merge.
+
+    Rotates the point cloud n_rotations times, runs inference each time,
+    un-rotates predictions back to world space, then merges all predictions
+    via vertex proximity deduplication.
+    """
+    assert pt_type == "colmap", "TTA only supports pt_type='colmap'"
+
+    points, colors, ade_colors, gestalt_colors = get_colmap_points(scene)
+    xyz_world = torch.from_numpy(points).float()
+    features = colors_to_features(gestalt_colors, ade_colors, rgb=colors, one_hot=False)
+
+    indices = torch.from_numpy(
+        sample_points(gestalt_colors, num_points, weight_vertex=100.0, weight_edge=10.0)
+    ).long()
+    xyz_sampled = xyz_world[indices]
+    features_sampled = features[indices]
+    centroid = xyz_sampled.mean(dim=0).numpy()
+
+    all_verts: list[np.ndarray] = []
+    all_edges: list[np.ndarray] = []
+    vert_offset = 0
+
+    for k in range(n_rotations):
+        theta = k * 2.0 * math.pi / n_rotations
+        R = _rotation_matrix_y(theta)
+
+        xyz_np = xyz_sampled.numpy()
+        xyz_rot_np = (xyz_np - centroid) @ R.T + centroid
+        xyz_rot = torch.from_numpy(xyz_rot_np).float()
+
+        min_xyz, max_xyz, q_diff = compute_norm_params(xyz_rot)
+        xyz_norm = (xyz_rot - min_xyz) / q_diff
+
+        results = model.inference(xyz_norm.to(device), features_sampled.to(device), threshold=score_threshold)
+        verts_norm = results["vertices"]
+        edges = results["edges"]
+
+        if len(edges) == 0:
+            continue
+
+        verts_rot_world = verts_norm * q_diff.numpy() + min_xyz.numpy()
+        verts_orig = (verts_rot_world - centroid) @ R.T + centroid  # R^T = R^{-1}
+
+        all_verts.append(verts_orig)
+        all_edges.append(edges + vert_offset)
+        vert_offset += len(verts_orig)
+
+    if not all_verts:
+        return np.zeros((0, 3), dtype=np.float32), np.zeros((0, 2), dtype=np.int64)
+
+    final_verts, final_edges = post_process_wireframe(
+        np.concatenate(all_verts, axis=0),
+        np.concatenate(all_edges, axis=0),
+        merge_distance_threshold=merge_distance_threshold,
+    )
+    return final_verts, final_edges