PCGao/MatchAnything

1"""23D visualization based on plotly.3Works for a small number of points and cameras, might be slow otherwise.4 51) Initialize a figure with `init_figure`62) Add 3D points, camera frustums, or both as a pycolmap.Reconstruction7 8Written by Paul-Edouard Sarlin and Philipp Lindenberger.9"""10 11from typing import Optional12 13import numpy as np14import plotly.graph_objects as go15import pycolmap16 17 18def to_homogeneous(points):19    pad = np.ones((points.shape[:-1] + (1,)), dtype=points.dtype)20    return np.concatenate([points, pad], axis=-1)21 22 23def init_figure(height: int = 800) -> go.Figure:24    """Initialize a 3D figure."""25    fig = go.Figure()26    axes = dict(27        visible=False,28        showbackground=False,29        showgrid=False,30        showline=False,31        showticklabels=True,32        autorange=True,33    )34    fig.update_layout(35        template="plotly_dark",36        height=height,37        scene_camera=dict(38            eye=dict(x=0.0, y=-0.1, z=-2),39            up=dict(x=0, y=-1.0, z=0),40            projection=dict(type="orthographic"),41        ),42        scene=dict(43            xaxis=axes,44            yaxis=axes,45            zaxis=axes,46            aspectmode="data",47            dragmode="orbit",48        ),49        margin=dict(l=0, r=0, b=0, t=0, pad=0),50        legend=dict(orientation="h", yanchor="top", y=0.99, xanchor="left", x=0.1),51    )52    return fig53 54 55def plot_points(56    fig: go.Figure,57    pts: np.ndarray,58    color: str = "rgba(255, 0, 0, 1)",59    ps: int = 2,60    colorscale: Optional[str] = None,61    name: Optional[str] = None,62):63    """Plot a set of 3D points."""64    x, y, z = pts.T65    tr = go.Scatter3d(66        x=x,67        y=y,68        z=z,69        mode="markers",70        name=name,71        legendgroup=name,72        marker=dict(size=ps, color=color, line_width=0.0, colorscale=colorscale),73    )74    fig.add_trace(tr)75 76 77def plot_camera(78    fig: go.Figure,79    R: np.ndarray,80    t: np.ndarray,81    K: np.ndarray,82    color: str = "rgb(0, 0, 255)",83    name: Optional[str] = None,84    legendgroup: Optional[str] = None,85    fill: bool = False,86    size: float = 1.0,87    text: Optional[str] = None,88):89    """Plot a camera frustum from pose and intrinsic matrix."""90    W, H = K[0, 2] * 2, K[1, 2] * 291    corners = np.array([[0, 0], [W, 0], [W, H], [0, H], [0, 0]])92    if size is not None:93        image_extent = max(size * W / 1024.0, size * H / 1024.0)94        world_extent = max(W, H) / (K[0, 0] + K[1, 1]) / 0.595        scale = 0.5 * image_extent / world_extent96    else:97        scale = 1.098    corners = to_homogeneous(corners) @ np.linalg.inv(K).T99    corners = (corners / 2 * scale) @ R.T + t100    legendgroup = legendgroup if legendgroup is not None else name101 102    x, y, z = np.concatenate(([t], corners)).T103    i = [0, 0, 0, 0]104    j = [1, 2, 3, 4]105    k = [2, 3, 4, 1]106 107    if fill:108        pyramid = go.Mesh3d(109            x=x,110            y=y,111            z=z,112            color=color,113            i=i,114            j=j,115            k=k,116            legendgroup=legendgroup,117            name=name,118            showlegend=False,119            hovertemplate=text.replace("\n", "<br>"),120        )121        fig.add_trace(pyramid)122 123    triangles = np.vstack((i, j, k)).T124    vertices = np.concatenate(([t], corners))125    tri_points = np.array([vertices[i] for i in triangles.reshape(-1)])126    x, y, z = tri_points.T127 128    pyramid = go.Scatter3d(129        x=x,130        y=y,131        z=z,132        mode="lines",133        legendgroup=legendgroup,134        name=name,135        line=dict(color=color, width=1),136        showlegend=False,137        hovertemplate=text.replace("\n", "<br>"),138    )139    fig.add_trace(pyramid)140 141 142def plot_camera_colmap(143    fig: go.Figure,144    image: pycolmap.Image,145    camera: pycolmap.Camera,146    name: Optional[str] = None,147    **kwargs,148):149    """Plot a camera frustum from PyCOLMAP objects"""150    world_t_camera = image.cam_from_world.inverse()151    plot_camera(152        fig,153        world_t_camera.rotation.matrix(),154        world_t_camera.translation,155        camera.calibration_matrix(),156        name=name or str(image.image_id),157        text=str(image),158        **kwargs,159    )160 161 162def plot_cameras(fig: go.Figure, reconstruction: pycolmap.Reconstruction, **kwargs):163    """Plot a camera as a cone with camera frustum."""164    for image_id, image in reconstruction.images.items():165        plot_camera_colmap(166            fig, image, reconstruction.cameras[image.camera_id], **kwargs167        )168 169 170def plot_reconstruction(171    fig: go.Figure,172    rec: pycolmap.Reconstruction,173    max_reproj_error: float = 6.0,174    color: str = "rgb(0, 0, 255)",175    name: Optional[str] = None,176    min_track_length: int = 2,177    points: bool = True,178    cameras: bool = True,179    points_rgb: bool = True,180    cs: float = 1.0,181):182    # Filter outliers183    bbs = rec.compute_bounding_box(0.001, 0.999)184    # Filter points, use original reproj error here185    p3Ds = [186        p3D187        for _, p3D in rec.points3D.items()188        if (189            (p3D.xyz >= bbs[0]).all()190            and (p3D.xyz <= bbs[1]).all()191            and p3D.error <= max_reproj_error192            and p3D.track.length() >= min_track_length193        )194    ]195    xyzs = [p3D.xyz for p3D in p3Ds]196    if points_rgb:197        pcolor = [p3D.color for p3D in p3Ds]198    else:199        pcolor = color200    if points:201        plot_points(fig, np.array(xyzs), color=pcolor, ps=1, name=name)202    if cameras:203        plot_cameras(fig, rec, color=color, legendgroup=name, size=cs)204