vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/gim/gluefactory/visualization/tools.py

@@ -0,0 +1,465 @@
+import inspect
+import sys
+import warnings
+
+import matplotlib.pyplot as plt
+import torch
+from matplotlib.backend_tools import ToolToggleBase
+from matplotlib.widgets import RadioButtons, Slider
+
+from ..geometry.epipolar import T_to_F, generalized_epi_dist
+from ..geometry.homography import sym_homography_error
+from ..visualization.viz2d import (
+    cm_ranking,
+    cm_RdGn,
+    draw_epipolar_line,
+    get_line,
+    plot_color_line_matches,
+    plot_heatmaps,
+    plot_keypoints,
+    plot_lines,
+    plot_matches,
+)
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    plt.rcParams["toolbar"] = "toolmanager"
+
+
+class RadioHideTool(ToolToggleBase):
+    """Show lines with a given gid."""
+
+    default_keymap = "R"
+    description = "Show by gid"
+    default_toggled = False
+    radio_group = "default"
+
+    def __init__(
+        self, *args, options=[], active=None, callback_fn=None, keymap="R", **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.f = 1.0
+        self.options = options
+        self.callback_fn = callback_fn
+        self.active = self.options.index(active) if active else 0
+        self.default_keymap = keymap
+
+        self.enabled = self.default_toggled
+
+    def build_radios(self):
+        w = 0.2
+        self.radios_ax = self.figure.add_axes([1.0 - w, 0.7, w, 0.2], zorder=1)
+        # self.radios_ax = self.figure.add_axes([0.5-w/2, 1.0-0.2, w, 0.2], zorder=1)
+        self.radios = RadioButtons(self.radios_ax, self.options, active=self.active)
+        self.radios.on_clicked(self.on_radio_clicked)
+
+    def enable(self, *args):
+        size = self.figure.get_size_inches()
+        size[0] *= self.f
+        self.build_radios()
+        self.figure.canvas.draw_idle()
+        self.enabled = True
+
+    def disable(self, *args):
+        size = self.figure.get_size_inches()
+        size[0] /= self.f
+        self.radios_ax.remove()
+        self.radios = None
+        self.figure.canvas.draw_idle()
+        self.enabled = False
+
+    def on_radio_clicked(self, value):
+        self.active = self.options.index(value)
+        enabled = self.enabled
+        if enabled:
+            self.disable()
+        if self.callback_fn is not None:
+            self.callback_fn(value)
+        if enabled:
+            self.enable()
+
+
+class ToggleTool(ToolToggleBase):
+    """Show lines with a given gid."""
+
+    default_keymap = "t"
+    description = "Show by gid"
+
+    def __init__(self, *args, callback_fn=None, keymap="t", **kwargs):
+        super().__init__(*args, **kwargs)
+        self.f = 1.0
+        self.callback_fn = callback_fn
+        self.default_keymap = keymap
+        self.enabled = self.default_toggled
+
+    def enable(self, *args):
+        self.callback_fn(True)
+
+    def disable(self, *args):
+        self.callback_fn(False)
+
+
+def add_whitespace_left(fig, factor):
+    w, h = fig.get_size_inches()
+    left = fig.subplotpars.left
+    fig.set_size_inches([w * (1 + factor), h])
+    fig.subplots_adjust(left=(factor + left) / (1 + factor))
+
+
+def add_whitespace_bottom(fig, factor):
+    w, h = fig.get_size_inches()
+    b = fig.subplotpars.bottom
+    fig.set_size_inches([w, h * (1 + factor)])
+    fig.subplots_adjust(bottom=(factor + b) / (1 + factor))
+    fig.canvas.draw_idle()
+
+
+class KeypointPlot:
+    plot_name = "keypoints"
+    required_keys = ["keypoints0", "keypoints1"]
+
+    def __init__(self, fig, axes, data, preds):
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            plot_keypoints([pred["keypoints0"][0], pred["keypoints1"][0]], axes=axes[i])
+
+
+class LinePlot:
+    plot_name = "lines"
+    required_keys = ["lines0", "lines1"]
+
+    def __init__(self, fig, axes, data, preds):
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            plot_lines([pred["lines0"][0], pred["lines1"][0]])
+
+
+class KeypointRankingPlot:
+    plot_name = "keypoint_ranking"
+    required_keys = ["keypoints0", "keypoints1", "keypoint_scores0", "keypoint_scores1"]
+
+    def __init__(self, fig, axes, data, preds):
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            kp0, kp1 = pred["keypoints0"][0], pred["keypoints1"][0]
+            sc0, sc1 = pred["keypoint_scores0"][0], pred["keypoint_scores1"][0]
+
+            plot_keypoints(
+                [kp0, kp1], axes=axes[i], colors=[cm_ranking(sc0), cm_ranking(sc1)]
+            )
+
+
+class KeypointScoresPlot:
+    plot_name = "keypoint_scores"
+    required_keys = ["keypoints0", "keypoints1", "keypoint_scores0", "keypoint_scores1"]
+
+    def __init__(self, fig, axes, data, preds):
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            kp0, kp1 = pred["keypoints0"][0], pred["keypoints1"][0]
+            sc0, sc1 = pred["keypoint_scores0"][0], pred["keypoint_scores1"][0]
+            plot_keypoints(
+                [kp0, kp1], axes=axes[i], colors=[cm_RdGn(sc0), cm_RdGn(sc1)]
+            )
+
+
+class HeatmapPlot:
+    plot_name = "heatmaps"
+    required_keys = ["heatmap0", "heatmap1"]
+
+    def __init__(self, fig, axes, data, preds):
+        self.artists = []
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            heatmaps = [pred["heatmap0"][0, 0], pred["heatmap1"][0, 0]]
+            heatmaps = [torch.sigmoid(h) if h.min() < 0.0 else h for h in heatmaps]
+            self.artists += plot_heatmaps(heatmaps, axes=axes[i], cmap="rainbow")
+
+    def clear(self):
+        for x in self.artists:
+            x.remove()
+
+
+class ImagePlot:
+    plot_name = "images"
+    required_keys = ["view0", "view1"]
+
+    def __init__(self, fig, axes, data, preds):
+        pass
+
+
+class MatchesPlot:
+    plot_name = "matches"
+    required_keys = ["keypoints0", "keypoints1", "matches0", "matching_scores0"]
+
+    def __init__(self, fig, axes, data, preds):
+        self.fig = fig
+        self.sbpars = {
+            k: v
+            for k, v in vars(fig.subplotpars).items()
+            if k in ["left", "right", "top", "bottom"]
+        }
+
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            plot_keypoints(
+                [pred["keypoints0"][0], pred["keypoints1"][0]],
+                axes=axes[i],
+                colors="blue",
+            )
+            kp0, kp1 = pred["keypoints0"][0], pred["keypoints1"][0]
+            m0 = pred["matches0"][0]
+            valid = m0 > -1
+            kpm0 = kp0[valid]
+            kpm1 = kp1[m0[valid]]
+            mscores = pred["matching_scores0"][0][valid]
+            plot_matches(
+                kpm0,
+                kpm1,
+                color=cm_RdGn(mscores).tolist(),
+                axes=axes[i],
+                labels=mscores,
+                lw=0.5,
+            )
+
+
+class LineMatchesPlot:
+    plot_name = "line_matches"
+    required_keys = ["lines0", "lines1", "line_matches0"]
+
+    def __init__(self, fig, axes, data, preds):
+        self.fig = fig
+        self.sbpars = {
+            k: v
+            for k, v in vars(fig.subplotpars).items()
+            if k in ["left", "right", "top", "bottom"]
+        }
+
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            lines0, lines1 = pred["lines0"][0], pred["lines1"][0]
+            m0 = pred["line_matches0"][0]
+            valid = m0 > -1
+            m_lines0 = lines0[valid]
+            m_lines1 = lines1[m0[valid]]
+            plot_color_line_matches([m_lines0, m_lines1])
+
+
+class GtMatchesPlot:
+    plot_name = "gt_matches"
+    required_keys = ["keypoints0", "keypoints1", "matches0", "gt_matches0"]
+
+    def __init__(self, fig, axes, data, preds):
+        self.fig = fig
+        self.sbpars = {
+            k: v
+            for k, v in vars(fig.subplotpars).items()
+            if k in ["left", "right", "top", "bottom"]
+        }
+
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            plot_keypoints(
+                [pred["keypoints0"][0], pred["keypoints1"][0]],
+                axes=axes[i],
+                colors="blue",
+            )
+            kp0, kp1 = pred["keypoints0"][0], pred["keypoints1"][0]
+            m0 = pred["matches0"][0]
+            gtm0 = pred["gt_matches0"][0]
+            valid = (m0 > -1) & (gtm0 >= -1)
+            kpm0 = kp0[valid]
+            kpm1 = kp1[m0[valid]]
+            correct = gtm0[valid] == m0[valid]
+            plot_matches(
+                kpm0,
+                kpm1,
+                color=cm_RdGn(correct).tolist(),
+                axes=axes[i],
+                labels=correct,
+                lw=0.5,
+            )
+
+
+class GtLineMatchesPlot:
+    plot_name = "gt_line_matches"
+    required_keys = ["lines0", "lines1", "line_matches0", "line_gt_matches0"]
+
+    def __init__(self, fig, axes, data, preds):
+        self.fig = fig
+        self.sbpars = {
+            k: v
+            for k, v in vars(fig.subplotpars).items()
+            if k in ["left", "right", "top", "bottom"]
+        }
+
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            lines0, lines1 = pred["lines0"][0], pred["lines1"][0]
+            m0 = pred["line_matches0"][0]
+            gtm0 = pred["gt_line_matches0"][0]
+            valid = (m0 > -1) & (gtm0 >= -1)
+            m_lines0 = lines0[valid]
+            m_lines1 = lines1[m0[valid]]
+            plot_color_line_matches([m_lines0, m_lines1])
+
+
+class HomographyMatchesPlot:
+    plot_name = "homography"
+    required_keys = ["keypoints0", "keypoints1", "matches0", "H_0to1"]
+
+    def __init__(self, fig, axes, data, preds):
+        self.fig = fig
+        self.sbpars = {
+            k: v
+            for k, v in vars(fig.subplotpars).items()
+            if k in ["left", "right", "top", "bottom"]
+        }
+
+        add_whitespace_bottom(fig, 0.1)
+
+        self.range_ax = fig.add_axes([0.3, 0.02, 0.4, 0.06])
+        self.range = Slider(
+            self.range_ax,
+            label="Homography Error",
+            valmin=0,
+            valmax=5,
+            valinit=3.0,
+            valstep=1.0,
+        )
+        self.range.on_changed(self.color_matches)
+
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            plot_keypoints(
+                [pred["keypoints0"][0], pred["keypoints1"][0]],
+                axes=axes[i],
+                colors="blue",
+            )
+            kp0, kp1 = pred["keypoints0"][0], pred["keypoints1"][0]
+            m0 = pred["matches0"][0]
+            valid = m0 > -1
+            kpm0 = kp0[valid]
+            kpm1 = kp1[m0[valid]]
+            errors = sym_homography_error(kpm0, kpm1, data["H_0to1"][0])
+            plot_matches(
+                kpm0,
+                kpm1,
+                color=cm_RdGn(errors < self.range.val).tolist(),
+                axes=axes[i],
+                labels=errors.numpy(),
+                lw=0.5,
+            )
+
+    def clear(self):
+        w, h = self.fig.get_size_inches()
+        self.fig.set_size_inches(w, h / 1.1)
+        self.fig.subplots_adjust(**self.sbpars)
+        self.range_ax.remove()
+
+    def color_matches(self, args):
+        for line in self.fig.artists:
+            label = line.get_label()
+            line.set_color(cm_RdGn([float(label) < args])[0])
+
+
+class EpipolarMatchesPlot:
+    plot_name = "epipolar_matches"
+    required_keys = ["keypoints0", "keypoints1", "matches0", "T_0to1", "view0", "view1"]
+
+    def __init__(self, fig, axes, data, preds):
+        self.fig = fig
+        self.axes = axes
+        self.sbpars = {
+            k: v
+            for k, v in vars(fig.subplotpars).items()
+            if k in ["left", "right", "top", "bottom"]
+        }
+
+        add_whitespace_bottom(fig, 0.1)
+
+        self.range_ax = fig.add_axes([0.3, 0.02, 0.4, 0.06])
+        self.range = Slider(
+            self.range_ax,
+            label="Epipolar Error [px]",
+            valmin=0,
+            valmax=5,
+            valinit=3.0,
+            valstep=1.0,
+        )
+        self.range.on_changed(self.color_matches)
+
+        camera0 = data["view0"]["camera"][0]
+        camera1 = data["view1"]["camera"][0]
+        T_0to1 = data["T_0to1"][0]
+
+        for i, name in enumerate(preds):
+            pred = preds[name]
+            plot_keypoints(
+                [pred["keypoints0"][0], pred["keypoints1"][0]],
+                axes=axes[i],
+                colors="blue",
+            )
+            kp0, kp1 = pred["keypoints0"][0], pred["keypoints1"][0]
+            m0 = pred["matches0"][0]
+            valid = m0 > -1
+            kpm0 = kp0[valid]
+            kpm1 = kp1[m0[valid]]
+
+            errors = generalized_epi_dist(
+                kpm0,
+                kpm1,
+                camera0,
+                camera1,
+                T_0to1,
+                all=False,
+                essential=False,
+            )
+            plot_matches(
+                kpm0,
+                kpm1,
+                color=cm_RdGn(errors < self.range.val).tolist(),
+                axes=axes[i],
+                labels=errors.numpy(),
+                lw=0.5,
+            )
+
+        self.F = T_to_F(camera0, camera1, T_0to1)
+
+    def clear(self):
+        w, h = self.fig.get_size_inches()
+        self.fig.set_size_inches(w, h / 1.1)
+        self.fig.subplots_adjust(**self.sbpars)
+        self.range_ax.remove()
+
+    def color_matches(self, args):
+        for art in self.fig.artists:
+            label = art.get_label()
+            if label is not None:
+                art.set_color(cm_RdGn([float(label) < args])[0])
+
+    def click_artist(self, event):
+        art = event.artist
+        if art.get_label() is not None:
+            if hasattr(art, "epilines"):
+                [
+                    x.set_visible(not x.get_visible())
+                    for x in art.epilines
+                    if x is not None
+                ]
+            else:
+                xy1 = art.xy1
+                xy2 = art.xy2
+                line0 = get_line(self.F.transpose(0, 1), xy2)[:, 0]
+                line1 = get_line(self.F, xy1)[:, 0]
+                art.epilines = [
+                    draw_epipolar_line(line0, art.axesA),
+                    draw_epipolar_line(line1, art.axesB),
+                ]
+
+
+__plot_dict__ = {
+    obj.plot_name: obj
+    for _, obj in inspect.getmembers(sys.modules[__name__], predicate=inspect.isclass)
+    if hasattr(obj, "plot_name")
+}

vismatch/third_party/gim/gluefactory/visualization/two_view_frame.py

@@ -0,0 +1,158 @@
+import pprint
+
+import numpy as np
+
+from . import viz2d
+from .tools import RadioHideTool, ToggleTool, __plot_dict__
+
+
+class FormatPrinter(pprint.PrettyPrinter):
+    def __init__(self, formats):
+        super(FormatPrinter, self).__init__()
+        self.formats = formats
+
+    def format(self, obj, ctx, maxlvl, lvl):
+        if type(obj) in self.formats:
+            return self.formats[type(obj)] % obj, 1, 0
+        return pprint.PrettyPrinter.format(self, obj, ctx, maxlvl, lvl)
+
+
+class TwoViewFrame:
+    default_conf = {
+        "default": "matches",
+        "summary_visible": False,
+    }
+
+    plot_dict = __plot_dict__
+
+    childs = []
+
+    event_to_image = [None, "color", "depth", "color+depth"]
+
+    def __init__(self, conf, data, preds, title=None, event=1, summaries=None):
+        self.conf = conf
+        self.data = data
+        self.preds = preds
+        self.names = list(preds.keys())
+        self.plot = self.event_to_image[event]
+        self.summaries = summaries
+        self.fig, self.axes, self.summary_arts = self.init_frame()
+        if title is not None:
+            self.fig.canvas.manager.set_window_title(title)
+
+        keys = None
+        for _, pred in preds.items():
+            if keys is None:
+                keys = set(pred.keys())
+            else:
+                keys = keys.intersection(pred.keys())
+        keys = keys.union(data.keys())
+
+        self.options = [
+            k for k, v in self.plot_dict.items() if set(v.required_keys).issubset(keys)
+        ]
+        self.handle = None
+        self.radios = self.fig.canvas.manager.toolmanager.add_tool(
+            "switch plot",
+            RadioHideTool,
+            options=self.options,
+            callback_fn=self.draw,
+            active=conf.default,
+            keymap="R",
+        )
+
+        self.toggle_summary = self.fig.canvas.manager.toolmanager.add_tool(
+            "toggle summary",
+            ToggleTool,
+            toggled=self.conf.summary_visible,
+            callback_fn=self.set_summary_visible,
+            keymap="t",
+        )
+
+        if self.fig.canvas.manager.toolbar is not None:
+            self.fig.canvas.manager.toolbar.add_tool("switch plot", "navigation")
+        self.draw(conf.default)
+
+    def init_frame(self):
+        """initialize frame"""
+        view0, view1 = self.data["view0"], self.data["view1"]
+        if self.plot == "color" or self.plot == "color+depth":
+            imgs = [
+                view0["image"][0].permute(1, 2, 0),
+                view1["image"][0].permute(1, 2, 0),
+            ]
+        elif self.plot == "depth":
+            imgs = [view0["depth"][0], view1["depth"][0]]
+        else:
+            raise ValueError(self.plot)
+        imgs = [imgs for _ in self.names]  # repeat for each model
+
+        fig, axes = viz2d.plot_image_grid(imgs, return_fig=True, titles=None, figs=5)
+        [viz2d.add_text(0, n, axes=axes[i]) for i, n in enumerate(self.names)]
+
+        if (
+            self.plot == "color+depth"
+            and "depth" in view0.keys()
+            and view0["depth"] is not None
+        ):
+            hmaps = [[view0["depth"][0], view1["depth"][0]] for _ in self.names]
+            [
+                viz2d.plot_heatmaps(hmaps[i], axes=axes[i], cmap="Spectral")
+                for i, _ in enumerate(hmaps)
+            ]
+
+        fig.canvas.mpl_connect("pick_event", self.click_artist)
+        if self.summaries is not None:
+            formatter = FormatPrinter({np.float32: "%.4f", np.float64: "%.4f"})
+            toggle_artists = [
+                viz2d.add_text(
+                    0,
+                    formatter.pformat(self.summaries[n]),
+                    axes=axes[i],
+                    pos=(0.01, 0.01),
+                    va="bottom",
+                    backgroundcolor=(0, 0, 0, 0.5),
+                    visible=self.conf.summary_visible,
+                )
+                for i, n in enumerate(self.names)
+            ]
+        else:
+            toggle_artists = []
+        return fig, axes, toggle_artists
+
+    def draw(self, value):
+        """redraw content in frame"""
+        self.clear()
+        self.conf.default = value
+        self.handle = self.plot_dict[value](self.fig, self.axes, self.data, self.preds)
+        return self.handle
+
+    def clear(self):
+        if self.handle is not None:
+            try:
+                self.handle.clear()
+            except AttributeError:
+                pass
+        self.handle = None
+        for row in self.axes:
+            for ax in row:
+                [li.remove() for li in ax.lines]
+                [c.remove() for c in ax.collections]
+        self.fig.artists.clear()
+        self.fig.canvas.draw_idle()
+        self.handle = None
+
+    def click_artist(self, event):
+        art = event.artist
+        select = art.get_arrowstyle().arrow == "-"
+        art.set_arrowstyle("<|-|>" if select else "-")
+        if select:
+            art.set_zorder(1)
+        if hasattr(self.handle, "click_artist"):
+            self.handle.click_artist(event)
+        self.fig.canvas.draw_idle()
+
+    def set_summary_visible(self, visible):
+        self.conf.summary_visible = visible
+        [s.set_visible(visible) for s in self.summary_arts]
+        self.fig.canvas.draw_idle()

vismatch/third_party/gim/gluefactory/visualization/visualize_batch.py

@@ -0,0 +1,57 @@
+import torch
+
+from ..utils.tensor import batch_to_device
+from .viz2d import cm_RdGn, plot_heatmaps, plot_image_grid, plot_keypoints, plot_matches
+
+
+def make_match_figures(pred_, data_, n_pairs=2):
+    # print first n pairs in batch
+    if "0to1" in pred_.keys():
+        pred_ = pred_["0to1"]
+    images, kpts, matches, mcolors = [], [], [], []
+    heatmaps = []
+    pred = batch_to_device(pred_, "cpu", non_blocking=False)
+    data = batch_to_device(data_, "cpu", non_blocking=False)
+
+    view0, view1 = data["view0"], data["view1"]
+
+    n_pairs = min(n_pairs, view0["image"].shape[0])
+    assert view0["image"].shape[0] >= n_pairs
+
+    kp0, kp1 = pred["keypoints0"], pred["keypoints1"]
+    m0 = pred["matches0"]
+    gtm0 = pred["gt_matches0"]
+
+    for i in range(n_pairs):
+        valid = (m0[i] > -1) & (gtm0[i] >= -1)
+        kpm0, kpm1 = kp0[i][valid].numpy(), kp1[i][m0[i][valid]].numpy()
+        images.append(
+            [view0["image"][i].permute(1, 2, 0), view1["image"][i].permute(1, 2, 0)]
+        )
+        kpts.append([kp0[i], kp1[i]])
+        matches.append((kpm0, kpm1))
+
+        correct = gtm0[i][valid] == m0[i][valid]
+
+        if "heatmap0" in pred.keys():
+            heatmaps.append(
+                [
+                    torch.sigmoid(pred["heatmap0"][i, 0]),
+                    torch.sigmoid(pred["heatmap1"][i, 0]),
+                ]
+            )
+        elif "depth" in view0.keys() and view0["depth"] is not None:
+            heatmaps.append([view0["depth"][i], view1["depth"][i]])
+
+        mcolors.append(cm_RdGn(correct).tolist())
+
+    fig, axes = plot_image_grid(images, return_fig=True, set_lim=True)
+    if len(heatmaps) > 0:
+        [plot_heatmaps(heatmaps[i], axes=axes[i], a=1.0) for i in range(n_pairs)]
+    [plot_keypoints(kpts[i], axes=axes[i], colors="royalblue") for i in range(n_pairs)]
+    [
+        plot_matches(*matches[i], color=mcolors[i], axes=axes[i], a=0.5, lw=1.0, ps=0.0)
+        for i in range(n_pairs)
+    ]
+
+    return {"matching": fig}