siclib/visualization/tools.py

import inspect
import sys
import warnings

import matplotlib.pyplot as plt
import numpy as np
import torch
from matplotlib.backend_tools import ToolToggleBase
from matplotlib.widgets import Button, RadioButtons

from siclib.geometry.camera import SimpleRadial as Camera
from siclib.geometry.gravity import Gravity
from siclib.geometry.perspective_fields import (
    get_latitude_field,
    get_perspective_field,
    get_up_field,
)
from siclib.models.utils.metrics import latitude_error, up_error
from siclib.utils.conversions import rad2deg
from siclib.visualization.viz2d import (
    add_text,
    plot_confidences,
    plot_heatmaps,
    plot_horizon_lines,
    plot_latitudes,
    plot_vector_fields,
)

# flake8: noqa
# mypy: ignore-errors

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.4, w, 0.5], 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 ImagePlot:
    plot_name = "image"
    required_keys = ["image"]

    def __init__(self, fig, axes, data, preds):
        pass


class HorizonLinePlot:
    plot_name = "horizon_line"
    required_keys = ["camera", "gravity"]

    def __init__(self, fig, axes, data, preds):
        for idx, name in enumerate(preds):
            pred = preds[name]
            gt_cam = data["camera"][0].detach().cpu()
            gt_gravity = data["gravity"][0].detach().cpu()
            plot_horizon_lines([gt_cam], [gt_gravity], line_colors="r", ax=[axes[0][idx]])

            if "camera" in pred and "gravity" in pred:
                pred_cam = Camera(pred["camera"][0].detach().cpu())
                gravity = Gravity(pred["gravity"][0].detach().cpu())
                plot_horizon_lines([pred_cam], [gravity], line_colors="yellow", ax=[axes[0][idx]])


class LatitudePlot:
    plot_name = "latitude"
    required_keys = ["latitude_field"]

    def __init__(self, fig, axes, data, preds):
        self.artists = []
        self.gt_mode = False  # Flag to track whether to display ground truth or predicted
        self.text_objects = []  # To store text objects

        self.fig = fig
        self.axes = axes
        self.data = data
        self.preds = preds

        # Create a toggle button on the lower left corner of the first axis
        self.ax_button = self.fig.add_axes([0.01, 0.02, 0.2, 0.06])
        self.button = Button(self.ax_button, "Toggle GT")
        self.button.on_clicked(self.toggle_display)

        self.update_plot()

    def toggle_display(self, event):
        # Toggle between ground truth and predicted latitudes
        self.gt_mode = not self.gt_mode
        self.update_plot()

    def update_plot(self):
        for x in self.artists:
            x.remove()
        for text in self.text_objects:
            text.remove()

        self.artists = []
        self.text_objects = []

        for idx, name in enumerate(self.preds):
            pred = self.preds[name]

            if self.gt_mode:
                latitude = self.data["latitude_field"][0][0]
                text = "\nGT"
            else:
                if "latitude_field" not in pred:
                    continue
                latitude = pred["latitude_field"][0][0]
                text = "\nPrediction"

            self.artists += plot_latitudes([latitude], axes=[self.axes[0][idx]])

            self.text_objects.append(add_text(idx, text))

        # Update the plot
        self.fig.canvas.draw()

    def clear(self):
        # Remove the button
        self.button.disconnect_events()
        self.ax_button.remove()

        for x in self.artists:
            x.remove()
        for text in self.text_objects:
            text.remove()

        self.artists = []
        self.text_objects = []


class LatitudeErrorPlot:
    plot_name = "latitude_error"
    required_keys = ["latitude_field"]

    def __init__(self, fig, axes, data, preds):
        self.artists = []
        for idx, name in enumerate(preds):
            pred = preds[name]
            gt = data["latitude_field"].detach().cpu()

            if "latitude_field" in pred:
                lat = pred["latitude_field"].detach().cpu()
                error = latitude_error(lat, gt)[0].numpy()

                if "latitude_confidence" in pred:
                    confidence = pred["latitude_confidence"].detach().cpu().numpy()
                    confidence = np.log10(confidence).clip(-5)
                    confidence = (confidence + 5) / (confidence.max() + 5)
                    arts = plot_heatmaps(
                        [error], cmap="turbo", axes=[axes[0][idx]], colorbar=True, a=confidence
                    )
                else:
                    arts = plot_heatmaps([error], cmap="turbo", axes=[axes[0][idx]], colorbar=True)
                self.artists += arts

    def clear(self):
        for x in self.artists:
            x.remove()
            x.colorbar.remove()

        self.artists = []


class LatitudeConfidencePlot:
    plot_name = "latitude_confidence"
    required_keys = []
    # required_keys = ["latitude_confidence"]

    def __init__(self, fig, axes, data, preds):
        self.artists = []
        for idx, name in enumerate(preds):
            pred = preds[name]

            if "latitude_confidence" in pred:
                arts = plot_confidences([pred["latitude_confidence"][0]], axes=[axes[0][idx]])
                self.artists += arts

    def clear(self):
        for x in self.artists:
            x.remove()
            x.colorbar.remove()

        self.artists = []


class UpPlot:
    plot_name = "up"
    required_keys = ["up_field"]

    def __init__(self, fig, axes, data, preds):
        self.artists = []
        self.gt_mode = False  # Flag to track whether to display ground truth or predicted
        self.text_objects = []  # To store text objects

        self.fig = fig
        self.axes = axes
        self.data = data
        self.preds = preds

        # Create a toggle button on the lower left corner of the first axis
        self.ax_button = self.fig.add_axes([0.01, 0.02, 0.2, 0.06])
        self.button = Button(self.ax_button, "Toggle GT")
        self.button.on_clicked(self.toggle_display)

        self.update_plot()

    def toggle_display(self, event):
        # Toggle between ground truth and predicted latitudes
        self.gt_mode = not self.gt_mode
        self.update_plot()

    def update_plot(self):
        for x in self.artists:
            x.remove()
        for text in self.text_objects:
            text.remove()

        self.artists = []
        self.text_objects = []

        for idx, name in enumerate(self.preds):
            pred = self.preds[name]

            if self.gt_mode:
                up = self.data["up_field"][0]
                text = "\nGT"
            else:
                if "up_field" not in pred:
                    continue
                up = pred["up_field"][0]
                text = "\nPrediction"

            # Plot up
            self.artists += plot_vector_fields([up], axes=[self.axes[0][idx]])

            self.text_objects.append(add_text(idx, text))

        # Update the plot
        self.fig.canvas.draw()

    def clear(self):
        # Remove the button
        self.button.disconnect_events()
        self.ax_button.remove()

        for x in self.artists:
            x.remove()
        for text in self.text_objects:
            text.remove()

        self.artists = []
        self.text_objects = []


class UpErrorPlot:
    plot_name = "up_error"
    required_keys = ["up_field"]

    def __init__(self, fig, axes, data, preds):
        self.artists = []
        for idx, name in enumerate(preds):
            pred = preds[name]
            gt = data["up_field"].detach().cpu()

            if "up_field" in pred:
                up = pred["up_field"].detach().cpu()
                error = up_error(up, gt)[0].numpy()

                if "up_confidence" in pred:
                    confidence = pred["up_confidence"].detach().cpu().numpy()
                    confidence = np.log10(confidence).clip(-5)
                    confidence = (confidence + 5) / (confidence.max() + 5)
                    arts = plot_heatmaps(
                        [error], cmap="turbo", axes=[axes[0][idx]], colorbar=True, a=confidence
                    )
                else:
                    arts = plot_heatmaps([error], cmap="turbo", axes=[axes[0][idx]], colorbar=True)
                self.artists += arts

    def clear(self):
        for x in self.artists:
            x.remove()
            x.colorbar.remove()

        self.artists = []


class UpConfidencePlot:
    plot_name = "up_confidence"
    required_keys = []
    # required_keys = ["up_confidence"]

    def __init__(self, fig, axes, data, preds):
        self.artists = []
        for idx, name in enumerate(preds):
            pred = preds[name]

            if "up_confidence" in pred:
                arts = plot_confidences([pred["up_confidence"][0]], axes=[axes[0][idx]])
                self.artists += arts

    def clear(self):
        for x in self.artists:
            x.remove()
            x.colorbar.remove()

        self.artists = []


class PerspectiveField:
    plot_name = "perspective_field"
    required_keys = ["camera", "gravity"]

    def __init__(self, fig, axes, data, preds):
        self.artists = []
        self.gt_mode = False  # Flag to track whether to display ground truth or predicted
        self.text_objects = []  # To store text objects

        self.fig = fig
        self.axes = axes
        self.data = data
        self.preds = preds

        # Create a toggle button on the lower left corner of the first axis
        self.ax_button = self.fig.add_axes([0.01, 0.02, 0.2, 0.06])
        self.button = Button(self.ax_button, "Toggle GT")
        self.button.on_clicked(self.toggle_display)

        self.update_plot()

    def toggle_display(self, event):
        # Toggle between ground truth and predicted latitudes
        self.gt_mode = not self.gt_mode
        self.update_plot()

    def update_plot(self):
        for x in self.artists:
            x.remove()
        for text in self.text_objects:
            text.remove()

        self.artists = []
        self.text_objects = []

        for idx, name in enumerate(self.preds):
            pred = self.preds[name]

            if self.gt_mode:
                camera = self.data["camera"]
                gravity = self.data["gravity"]
                text = "\nGT"
            else:
                camera = pred["camera"]
                gravity = pred["gravity"]
                text = "\nPrediction"
                camera = Camera(camera)
                gravity = Gravity(gravity)

            up, latitude = get_perspective_field(camera, gravity)

            self.artists += plot_latitudes([latitude[0, 0]], axes=[self.axes[0][idx]])
            self.artists += plot_vector_fields([up[0]], axes=[self.axes[0][idx]])

            self.text_objects.append(add_text(idx, text))

        # Update the plot
        self.fig.canvas.draw()

    def clear(self):
        # Remove the button
        self.button.disconnect_events()
        self.ax_button.remove()

        for x in self.artists:
            x.remove()
        for text in self.text_objects:
            text.remove()

        self.artists = []
        self.text_objects = []


__plot_dict__ = {
    obj.plot_name: obj
    for _, obj in inspect.getmembers(sys.modules[__name__], predicate=inspect.isclass)
    if hasattr(obj, "plot_name")
}