PCGao/MatchAnything

1"""22D visualization primitives based on Matplotlib.3 41) Plot images with `plot_images`.52) Call `plot_keypoints` or `plot_matches` any number of times.63) Optionally: save a .png or .pdf plot (nice in papers!) with `save_plot`.7"""8 9import matplotlib10import matplotlib.pyplot as plt11import matplotlib.patheffects as path_effects12import numpy as np13 14 15def cm_RdGn(x):16    """Custom colormap: red (0) -> yellow (0.5) -> green (1)."""17    x = np.clip(x, 0, 1)[..., None] * 218    c = x * np.array([[0, 1.0, 0]]) + (2 - x) * np.array([[1.0, 0, 0]])19    return np.clip(c, 0, 1)20 21 22def plot_images(23    imgs, titles=None, cmaps="gray", dpi=100, pad=0.5, adaptive=True, figsize=4.524):25    """Plot a set of images horizontally.26    Args:27        imgs: a list of NumPy or PyTorch images, RGB (H, W, 3) or mono (H, W).28        titles: a list of strings, as titles for each image.29        cmaps: colormaps for monochrome images.30        adaptive: whether the figure size should fit the image aspect ratios.31    """32    n = len(imgs)33    if not isinstance(cmaps, (list, tuple)):34        cmaps = [cmaps] * n35 36    if adaptive:37        ratios = [i.shape[1] / i.shape[0] for i in imgs]  # W / H38    else:39        ratios = [4 / 3] * n40    figsize = [sum(ratios) * figsize, figsize]41    fig, axs = plt.subplots(42        1, n, figsize=figsize, dpi=dpi, gridspec_kw={"width_ratios": ratios}43    )44    if n == 1:45        axs = [axs]46    for i, (img, ax) in enumerate(zip(imgs, axs)):47        ax.imshow(img, cmap=plt.get_cmap(cmaps[i]))48        ax.set_axis_off()49        if titles:50            ax.set_title(titles[i])51    fig.tight_layout(pad=pad)52    return fig53 54 55# def plot_keypoints(kpts, colors="lime", ps=4):56def plot_keypoints(kpts, colors=(0.8392925797770089, 0.1845444059976932, 0.1528642829680892), ps=4):57    """Plot keypoints for existing images.58    Args:59        kpts: list of ndarrays of size (N, 2).60        colors: string, or list of list of tuples (one for each keypoints).61        ps: size of the keypoints as float.62    """63    if not isinstance(colors, list):64        colors = [colors] * len(kpts)65    axes = plt.gcf().axes66    try:67        for a, k, c in zip(axes, kpts, colors):68            a.scatter(k[:, 0], k[:, 1], c=c, s=ps, linewidths=0)69    except IndexError:70        pass71 72 73def plot_matches(kpts0, kpts1, color=None, lw=1.5, ps=4, indices=(0, 1), a=1.0):74    """Plot matches for a pair of existing images.75    Args:76        kpts0, kpts1: corresponding keypoints of size (N, 2).77        color: color of each match, string or RGB tuple. Random if not given.78        lw: width of the lines.79        ps: size of the end points (no endpoint if ps=0)80        indices: indices of the images to draw the matches on.81        a: alpha opacity of the match lines.82    """83    fig = plt.gcf()84    ax = fig.axes85    assert len(ax) > max(indices)86    ax0, ax1 = ax[indices[0]], ax[indices[1]]87    fig.canvas.draw()88 89    assert len(kpts0) == len(kpts1)90    if color is None:91        color = matplotlib.cm.hsv(np.random.rand(len(kpts0))).tolist()92    elif len(color) > 0 and not isinstance(color[0], (tuple, list)):93        color = [color] * len(kpts0)94 95    if lw > 0:96        # transform the points into the figure coordinate system97        for i in range(len(kpts0)):98            fig.add_artist(99                matplotlib.patches.ConnectionPatch(100                    xyA=(kpts0[i, 0], kpts0[i, 1]),101                    coordsA=ax0.transData,102                    xyB=(kpts1[i, 0], kpts1[i, 1]),103                    coordsB=ax1.transData,104                    zorder=1,105                    color=color[i],106                    linewidth=lw,107                    alpha=a,108                )109            )110 111    # freeze the axes to prevent the transform to change112    ax0.autoscale(enable=False)113    ax1.autoscale(enable=False)114 115    if ps > 0:116        ax0.scatter(kpts0[:, 0], kpts0[:, 1], c=color, s=ps)117        ax1.scatter(kpts1[:, 0], kpts1[:, 1], c=color, s=ps)118 119 120def add_text(121    idx,122    text,123    pos=(0.01, 0.99),124    fs=15,125    color="w",126    lcolor="k",127    lwidth=2,128    ha="left",129    va="top",130):131    ax = plt.gcf().axes[idx]132    t = ax.text(133        *pos, text, fontsize=fs, ha=ha, va=va, color=color, transform=ax.transAxes134    )135    if lcolor is not None:136        t.set_path_effects(137            [138                path_effects.Stroke(linewidth=lwidth, foreground=lcolor),139                path_effects.Normal(),140            ]141        )142 143 144def save_plot(path, **kw):145    """Save the current figure without any white margin."""146    plt.savefig(path, bbox_inches="tight", pad_inches=0, **kw)147