|
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": "image", |
|
"summary_visible": False, |
|
} |
|
|
|
plot_dict = __plot_dict__ |
|
|
|
childs = [] |
|
|
|
event_to_image = [None, "image", "horizon_line", "lat_pred", "lat_gt"] |
|
|
|
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(): |
|
keys = set(pred.keys()) if keys is None else 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""" |
|
imgs = [[self.data["image"][0].permute(1, 2, 0) for _ in self.names]] |
|
|
|
|
|
fig, axes = viz2d.plot_image_grid(imgs, return_fig=True, titles=None, figs=5) |
|
[viz2d.add_text(i, n, axes=axes[0]) for i, n in enumerate(self.names)] |
|
|
|
fig.canvas.mpl_connect("pick_event", self.click_artist) |
|
if self.summaries is not None: |
|
font_size = 7 |
|
formatter = FormatPrinter({np.float32: "%.4f", np.float64: "%.4f"}) |
|
toggle_artists = [ |
|
viz2d.add_text( |
|
i, |
|
formatter.pformat(self.summaries[n]), |
|
axes=axes[0], |
|
pos=(0.01, 0.01), |
|
va="bottom", |
|
backgroundcolor=(0, 0, 0, 0.5), |
|
visible=self.conf.summary_visible, |
|
fs=font_size, |
|
) |
|
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() |
|
|
