def task(args):
task_id, a, slice_ = args
width, height = a.shape
b = io.BytesIO()
x = np.arange(width)
y = np.arange(height)
X, Y = np.meshgrid(x, y)
fig, ax = plt.subplots(figsize=(width / 100, height / 100), frameon=False, dpi=100)
fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
quad_mesh = ax.pcolormesh(X, Y, a.T, cmap='viridis')
ax.grid(False)
ax.axis('off')
# ax.semilogy()
rect = Rectangle((0, 0), 0, height, linewidth=0, facecolor='black', alpha=0.5)
ax.add_patch(rect)
for i in range(slice_.start, slice_.stop):
rect.set_bounds((0, 0, i, height))
fig.savefig(b, format='rgba', dpi=100)
print('task_id:', task_id, 'i', i)
return task_id, b.getvalue()
class CustomToolbar(NavToolbar):
toolitems = NavToolbar.toolitems + (
(None, None, None, None),
("ROI", "Select ROI", "selection", "_on_custom_select"),
)
def __init__(self, plotCanvas):
# create the default toolbar
NavToolbar.__init__(self, plotCanvas)
self.selector = RectSelector(
self.canvas.figure.axes[0], self.onSelect, button=[1, 3], minspanx=5, minspany=5 # don't use middle button
)
self.selector.set_active(True)
self.ax = self.canvas.figure.axes[0]
self.roi = None
self.fixedSize = False
if wx.Platform == "__WXMAC__":
self.to_draw = Rectangle(
(0, 0), 0, 1, visible=False, facecolor="yellow", edgecolor="black", alpha=0.5, fill=True
)
self.ax.add_patch(self.to_draw)
self.background = None
def _init_toolbar(self):
self._parent = self.canvas.GetParent()
self.wx_ids = {}
for text, tooltip_text, image_file, callback in self.toolitems:
if text is None:
self.AddSeparator()
continue
self.wx_ids[text] = wx.NewId()
try:
bitmap = _load_bitmap(image_file + ".png")
except IOError:
bitmap = wx.Bitmap(image_file + ".png")
if text in ["Pan", "Zoom", "ROI"]:
self.AddCheckTool(self.wx_ids[text], bitmap, shortHelp=text, longHelp=tooltip_text)
else:
self.AddSimpleTool(self.wx_ids[text], bitmap, text, tooltip_text)
bind(self, wx.EVT_TOOL, getattr(self, callback), id=self.wx_ids[text])
self.ToggleTool(self.wx_ids["ROI"], True)
self.Realize()
def _set_markers(self):
self.canvas.parentFrame.set_markers()
def _update_view(self):
NavToolbar._update_view(self)
self._set_markers()
# MacOS needs a forced draw to update plot
if wx.Platform == "__WXMAC__":
self.canvas.draw()
def draw(self):
self._set_markers()
NavToolbar.draw(self)
# MacOS needs a forced draw to update plot
if wx.Platform == "__WXMAC__":
self.canvas.draw()
def zoom(self, ev):
if wx.Platform == "__WXMAC__":
self.ToggleTool(self.wx_ids["Zoom"], self.GetToolState(self.wx_ids["Zoom"]))
NavToolbar.zoom(self, ev)
def pan(self, ev):
if wx.Platform == "__WXMAC__":
self.ToggleTool(self.wx_ids["Pan"], self.GetToolState(self.wx_ids["Pan"]))
NavToolbar.pan(self, ev)
def press_zoom(self, ev):
if wx.Platform == "__WXMAC__":
self.update_background()
self.to_draw.set_visible(True)
NavToolbar.press_zoom(self, ev)
def release_zoom(self, ev):
if wx.Platform == "__WXMAC__":
self.to_draw.set_visible(False)
NavToolbar.release_zoom(self, ev)
def draw_rubberband(self, event, x0, y0, x1, y1):
# XOR does not work on MacOS ...
if wx.Platform != "__WXMAC__":
NavToolbar.draw_rubberband(self, event, x0, y0, x1, y1)
else:
if self.background is not None:
self.canvas.restore_region(self.background)
c0, c1 = self.ax.transData.inverted().transform([[x0, y0], [x1, y1]])
l, b = c0
r, t = c1
self.to_draw.set_bounds(l, b, r - l, t - b)
self.ax.draw_artist(self.to_draw)
self.canvas.blit(self.ax.bbox)
def update_background(self):
"""force an update of the background"""
self.background = self.canvas.copy_from_bbox(self.ax.bbox)
# Turn on selection
# TODO: Proper handling of states, actual functionality.
def _on_custom_select(self, evt):
# for id in ['Zoom','Pan']:
# self.ToggleTool(self.wx_ids[id], False)
# print('Select ROI: %s' % (self.GetToolState(self.wx_ids['ROI'])))
# self.ToggleTool(self.wx_ids['ROI'],
# self.GetToolState(self.wx_ids['ROI']) )
self.toggle_selector()
# print('Select ROI: %s' % (self.GetToolState(self.wx_ids['ROI'])))
def onSelect(self, eclick, erelease):
"eclick and erelease are matplotlib events at press and release"
# print(' startposition : (%f, %f)' % (eclick.xdata, eclick.ydata))
# print(' endposition : (%f, %f)' % (erelease.xdata, erelease.ydata))
# print(' used button : ', eclick.button)
self.updateROI(
min(eclick.xdata, erelease.xdata),
min(eclick.ydata, erelease.ydata),
abs(eclick.xdata - erelease.xdata),
abs(eclick.ydata - erelease.ydata),
)
if self.canvas.parentFrame.fixedNumberCB.IsChecked():
# We are working in the fixed-number mode
# We need to find new roi for this center point
# The handler will call the update ROI function for us.
self.canvas.parentFrame.handleROIforN()
def updateROI(self, x, y, w, h):
if self.roi is None:
# print('upd ROI:', x, y, w, h)
self.roi = Rectangle((x, y), w, h, ls="solid", lw=2, color="r", fill=False, zorder=5)
self.canvas.figure.axes[0].add_patch(self.roi)
else:
self.roi.set_bounds(x, y, w, h)
self.updateCanvas()
def toggle_selector(self):
self.selector.set_active(not self.selector.active)
def onFixedSize(self, ev):
self.fixedSize = ev.IsChecked()
self.updateCanvas()
def onWidthChange(self, ev):
if self.roi:
x = self.roi.get_x()
w = self.roi.get_width()
nw = ev.GetValue()
dw = {"C": (w - nw) / 2, "L": 0, "R": w - nw}[self.canvas.parentFrame.anchorRB.GetStringSelection()[0]]
self.roi.set_x(x + dw)
self.roi.set_width(nw)
self.updateCanvas()
def onHeightChange(self, ev):
if self.roi:
y = self.roi.get_y()
h = self.roi.get_height()
nh = ev.GetValue()
dh = {"C": (h - nh) / 2, "B": 0, "T": h - nh}[self.canvas.parentFrame.anchorRB.GetStringSelection()[-1]]
self.roi.set_y(y + dh)
self.roi.set_height(nh)
self.updateCanvas()
def updateCanvas(self, redraw=True):
if self.roi:
self.canvas.parentFrame.showROI(
self.roi.get_x(), self.roi.get_y(), self.roi.get_width(), self.roi.get_height()
)
self.canvas.parentFrame.setWH(self.roi.get_width(), self.roi.get_height())
if self.fixedSize:
self.selector.setSize(self.roi.get_width(), self.roi.get_height())
else:
self.selector.setSize()
if redraw:
self.draw()
class Alignment(Axes):
"""
matplotlib.axes.Axes subclass for rendering sequence alignments.
"""
def __init__(self, fig, rect, *args, **kwargs):
self.aln = kwargs.pop("aln")
nrows = len(self.aln)
ncols = self.aln.get_alignment_length()
self.alnidx = numpy.arange(ncols)
self.app = kwargs.pop("app", None)
self.showy = kwargs.pop('showy', True)
Axes.__init__(self, fig, rect, *args, **kwargs)
rgb = mpl_colors.colorConverter.to_rgb
gray = rgb('gray')
d = defaultdict(lambda:gray)
d["A"] = rgb("red")
d["a"] = rgb("red")
d["C"] = rgb("blue")
d["c"] = rgb("blue")
d["G"] = rgb("green")
d["g"] = rgb("green")
d["T"] = rgb("yellow")
d["t"] = rgb("yellow")
self.cmap = d
self.selector = RectangleSelector(
self, self.select_rectangle, useblit=True
)
def f(e):
if e.button != 1: return True
else: return RectangleSelector.ignore(self.selector, e)
self.selector.ignore = f
self.selected_rectangle = Rectangle(
[0,0],0,0, facecolor='white', edgecolor='cyan', alpha=0.3
)
self.add_patch(self.selected_rectangle)
self.highlight_find_collection = None
def plot_aln(self):
cmap = self.cmap
self.ntax = len(self.aln); self.nchar = self.aln.get_alignment_length()
a = numpy.array([ [ cmap[base] for base in x.seq ]
for x in self.aln ])
self.array = a
self.imshow(a, interpolation='nearest', aspect='auto', origin='lower')
y = [ i+0.5 for i in xrange(self.ntax) ]
labels = [ x.id for x in self.aln ]
## locator.bin_boundaries(1,ntax)
## locator.view_limits(1,ntax)
if self.showy:
locator = MaxNLocator(nbins=50, integer=True)
self.yaxis.set_major_locator(locator)
def fmt(x, pos=None):
if x<0: return ""
try: return labels[int(round(x))]
except: pass
return ""
self.yaxis.set_major_formatter(FuncFormatter(fmt))
else:
self.yaxis.set_major_locator(NullLocator())
return self
def select_rectangle(self, e0, e1):
x0, x1 = map(int, sorted((e0.xdata+0.5, e1.xdata+0.5)))
y0, y1 = map(int, sorted((e0.ydata+0.5, e1.ydata+0.5)))
self.selected_chars = (x0, x1)
self.selected_taxa = (y0, y1)
self.selected_rectangle.set_bounds(x0-0.5,y0-0.5,x1-x0+1,y1-y0+1)
self.app.figure.canvas.draw_idle()
def highlight_find(self, substr):
if not substr:
if self.highlight_find_collection:
self.highlight_find_collection.remove()
self.highlight_find_collection = None
return
N = len(substr)
v = []
for y, x in align.find(self.aln, substr):
r = Rectangle(
[x-0.5,y-0.5], N, 1,
facecolor='cyan', edgecolor='cyan', alpha=0.7
)
v.append(r)
if self.highlight_find_collection:
self.highlight_find_collection.remove()
c = PatchCollection(v, True)
self.highlight_find_collection = self.add_collection(c)
self.app.figure.canvas.draw_idle()
def extract_selected(self):
r0, r1 = self.selected_taxa
c0, c1 = self.selected_chars
return self.aln[r0:r1+1,c0:c1+1]
def zoom_cxy(self, x=0.1, y=0.1, cx=None, cy=None):
"""
Zoom the x and y axes in by the specified proportion of the
current view, with a fixed data point (cx, cy)
"""
transform = self.transData.inverted().transform
xlim = self.get_xlim(); xmid = sum(xlim)*0.5
ylim = self.get_ylim(); ymid = sum(ylim)*0.5
bb = self.get_window_extent()
bbx = bb.expanded(1.0-x,1.0-y)
points = transform(bbx.get_points())
x0, x1 = points[:,0]; y0, y1 = points[:,1]
deltax = xmid-x0; deltay = ymid-y0
cx = cx or xmid; cy = cy or ymid
xoff = (cx-xmid)*x
self.set_xlim(xmid-deltax+xoff, xmid+deltax+xoff)
yoff = (cy-ymid)*y
self.set_ylim(ymid-deltay+yoff, ymid+deltay+yoff)
def zoom(self, x=0.1, y=0.1, cx=None, cy=None):
"""
Zoom the x and y axes in by the specified proportion of the
current view.
"""
# get the function to convert display coordinates to data
# coordinates
transform = self.transData.inverted().transform
xlim = self.get_xlim()
ylim = self.get_ylim()
bb = self.get_window_extent()
bbx = bb.expanded(1.0-x,1.0-y)
points = transform(bbx.get_points())
x0, x1 = points[:,0]; y0, y1 = points[:,1]
deltax = x0 - xlim[0]; deltay = y0 - ylim[0]
self.set_xlim(xlim[0]+deltax, xlim[1]-deltax)
self.set_ylim(ylim[0]+deltay, ylim[1]-deltay)
def center_y(self, y):
ymin, ymax = self.get_ylim()
yoff = (ymax - ymin) * 0.5
self.set_ylim(y-yoff, y+yoff)
def center_x(self, x, offset=0.3):
xmin, xmax = self.get_xlim()
xspan = xmax - xmin
xoff = xspan*0.5 + xspan*offset
self.set_xlim(x-xoff, x+xoff)
def scroll(self, x, y):
x0, x1 = self.get_xlim()
y0, y1 = self.get_ylim()
xd = (x1-x0)*x
yd = (y1-y0)*y
self.set_xlim(x0+xd, x1+xd)
self.set_ylim(y0+yd, y1+yd)
def home(self):
self.set_xlim(0, self.nchar)
self.set_ylim(self.ntax, 0)
class Alignment(Axes):
"""
matplotlib.axes.Axes subclass for rendering sequence alignments.
"""
def __init__(self, fig, rect, *args, **kwargs):
self.aln = kwargs.pop("aln")
nrows = len(self.aln)
ncols = self.aln.get_alignment_length()
self.alnidx = numpy.arange(ncols)
self.app = kwargs.pop("app", None)
self.showy = kwargs.pop('showy', True)
Axes.__init__(self, fig, rect, *args, **kwargs)
rgb = mpl_colors.colorConverter.to_rgb
gray = rgb('gray')
d = defaultdict(lambda: gray)
d["A"] = rgb("red")
d["a"] = rgb("red")
d["C"] = rgb("blue")
d["c"] = rgb("blue")
d["G"] = rgb("green")
d["g"] = rgb("green")
d["T"] = rgb("yellow")
d["t"] = rgb("yellow")
self.cmap = d
self.selector = RectangleSelector(self,
self.select_rectangle,
useblit=True)
def f(e):
if e.button != 1: return True
else: return RectangleSelector.ignore(self.selector, e)
self.selector.ignore = f
self.selected_rectangle = Rectangle([0, 0],
0,
0,
facecolor='white',
edgecolor='cyan',
alpha=0.3)
self.add_patch(self.selected_rectangle)
self.highlight_find_collection = None
def plot_aln(self):
cmap = self.cmap
self.ntax = len(self.aln)
self.nchar = self.aln.get_alignment_length()
a = numpy.array([[cmap[base] for base in x.seq] for x in self.aln])
self.array = a
self.imshow(a, interpolation='nearest', aspect='auto', origin='lower')
y = [i + 0.5 for i in range(self.ntax)]
labels = [x.id for x in self.aln]
## locator.bin_boundaries(1,ntax)
## locator.view_limits(1,ntax)
if self.showy:
locator = MaxNLocator(nbins=50, integer=True)
self.yaxis.set_major_locator(locator)
def fmt(x, pos=None):
if x < 0: return ""
try:
return labels[int(round(x))]
except:
pass
return ""
self.yaxis.set_major_formatter(FuncFormatter(fmt))
else:
self.yaxis.set_major_locator(NullLocator())
return self
def select_rectangle(self, e0, e1):
x0, x1 = list(map(int, sorted((e0.xdata + 0.5, e1.xdata + 0.5))))
y0, y1 = list(map(int, sorted((e0.ydata + 0.5, e1.ydata + 0.5))))
self.selected_chars = (x0, x1)
self.selected_taxa = (y0, y1)
self.selected_rectangle.set_bounds(x0 - 0.5, y0 - 0.5, x1 - x0 + 1,
y1 - y0 + 1)
self.app.figure.canvas.draw_idle()
def highlight_find(self, substr):
if not substr:
if self.highlight_find_collection:
self.highlight_find_collection.remove()
self.highlight_find_collection = None
return
N = len(substr)
v = []
for y, x in align.find(self.aln, substr):
r = Rectangle([x - 0.5, y - 0.5],
N,
1,
facecolor='cyan',
edgecolor='cyan',
alpha=0.7)
v.append(r)
if self.highlight_find_collection:
self.highlight_find_collection.remove()
c = PatchCollection(v, True)
self.highlight_find_collection = self.add_collection(c)
self.app.figure.canvas.draw_idle()
def extract_selected(self):
r0, r1 = self.selected_taxa
c0, c1 = self.selected_chars
return self.aln[r0:r1 + 1, c0:c1 + 1]
def zoom_cxy(self, x=0.1, y=0.1, cx=None, cy=None):
"""
Zoom the x and y axes in by the specified proportion of the
current view, with a fixed data point (cx, cy)
"""
transform = self.transData.inverted().transform
xlim = self.get_xlim()
xmid = sum(xlim) * 0.5
ylim = self.get_ylim()
ymid = sum(ylim) * 0.5
bb = self.get_window_extent()
bbx = bb.expanded(1.0 - x, 1.0 - y)
points = transform(bbx.get_points())
x0, x1 = points[:, 0]
y0, y1 = points[:, 1]
deltax = xmid - x0
deltay = ymid - y0
cx = cx or xmid
cy = cy or ymid
xoff = (cx - xmid) * x
self.set_xlim(xmid - deltax + xoff, xmid + deltax + xoff)
yoff = (cy - ymid) * y
self.set_ylim(ymid - deltay + yoff, ymid + deltay + yoff)
def zoom(self, x=0.1, y=0.1, cx=None, cy=None):
"""
Zoom the x and y axes in by the specified proportion of the
current view.
"""
# get the function to convert display coordinates to data
# coordinates
transform = self.transData.inverted().transform
xlim = self.get_xlim()
ylim = self.get_ylim()
bb = self.get_window_extent()
bbx = bb.expanded(1.0 - x, 1.0 - y)
points = transform(bbx.get_points())
x0, x1 = points[:, 0]
y0, y1 = points[:, 1]
deltax = x0 - xlim[0]
deltay = y0 - ylim[0]
self.set_xlim(xlim[0] + deltax, xlim[1] - deltax)
self.set_ylim(ylim[0] + deltay, ylim[1] - deltay)
def center_y(self, y):
ymin, ymax = self.get_ylim()
yoff = (ymax - ymin) * 0.5
self.set_ylim(y - yoff, y + yoff)
def center_x(self, x, offset=0.3):
xmin, xmax = self.get_xlim()
xspan = xmax - xmin
xoff = xspan * 0.5 + xspan * offset
self.set_xlim(x - xoff, x + xoff)
def scroll(self, x, y):
x0, x1 = self.get_xlim()
y0, y1 = self.get_ylim()
xd = (x1 - x0) * x
yd = (y1 - y0) * y
self.set_xlim(x0 + xd, x1 + xd)
self.set_ylim(y0 + yd, y1 + yd)
def home(self):
self.set_xlim(0, self.nchar)
self.set_ylim(self.ntax, 0)
class StatsPanel(wx.Panel):
def __init__(self, parent):
wx.Panel.__init__(self, parent, wx.ID_ANY, wx.DefaultPosition, wx.DefaultSize)
self.ztv_frame = self.GetTopLevelParent()
self.ztv_frame.primary_image_panel.popup_menu_cursor_modes.append('Stats box')
self.ztv_frame.primary_image_panel.available_cursor_modes['Stats box'] = {
'set-to-mode':self.set_cursor_to_stats_box_mode,
'on_button_press':self.on_button_press,
'on_motion':self.on_motion,
'on_button_release':self.on_button_release}
self.textentry_font = wx.Font(14, wx.FONTFAMILY_MODERN, wx.NORMAL, wx.FONTWEIGHT_LIGHT, False)
self.stats_info = None
self.last_string_values = {'x0':'', 'xsize':'', 'x1':'', 'y0':'', 'ysize':'', 'y1':''}
self.stats_rect = Rectangle((0, 0), 10, 10, color='magenta', fill=False, zorder=100)
# use self.stats_rect as where we store/retrieve the x0,y0,x1,y1
# x0,y0,x1,y1 should be limited to range of 0 to shape-1
# but, stats should be calculated over e.g. x0:x1+1 (so that have pixels to do stats on even if x0==x1)
# and, width/height of stats_rect should always be >= 0
values_sizer = wx.FlexGridSizer( 10, 5, 0, 0 )
values_sizer.SetFlexibleDirection( wx.BOTH )
values_sizer.SetNonFlexibleGrowMode( wx.FLEX_GROWMODE_SPECIFIED )
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
self.low_static_text = wx.StaticText( self, wx.ID_ANY, u"Low", wx.DefaultPosition, wx.DefaultSize, wx.ALIGN_RIGHT )
self.low_static_text.Wrap( -1 )
values_sizer.Add(self.low_static_text, 0, wx.ALL|wx.ALIGN_CENTER_HORIZONTAL, 0)
self.low_static_text = wx.StaticText( self, wx.ID_ANY, u"# pix", wx.DefaultPosition, wx.DefaultSize, 0 )
self.low_static_text.Wrap( -1 )
values_sizer.Add(self.low_static_text, 0, wx.ALL|wx.ALIGN_CENTER_HORIZONTAL, 0)
self.high_static_text = wx.StaticText( self, wx.ID_ANY, u"High", wx.DefaultPosition, wx.DefaultSize, 0 )
self.high_static_text.Wrap( -1 )
values_sizer.Add(self.high_static_text, 0, wx.ALL|wx.ALIGN_CENTER_HORIZONTAL, 0)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
self.x_static_text = wx.StaticText( self, wx.ID_ANY, u"x", wx.DefaultPosition, wx.DefaultSize, 0 )
self.x_static_text.Wrap( -1 )
values_sizer.Add(self.x_static_text, 0, wx.ALL|wx.ALIGN_CENTER_VERTICAL, 0)
self.x0_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.x0_textctrl.SetFont(self.textentry_font)
values_sizer.Add(self.x0_textctrl, 0, wx.ALL, 2)
self.x0_textctrl.Bind(wx.EVT_TEXT, self.x0_textctrl_changed)
self.x0_textctrl.Bind(wx.EVT_TEXT_ENTER, self.x0_textctrl_entered)
self.xsize_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.xsize_textctrl.SetFont(self.textentry_font)
values_sizer.Add(self.xsize_textctrl, 0, wx.ALL, 2)
self.xsize_textctrl.Bind(wx.EVT_TEXT, self.xsize_textctrl_changed)
self.xsize_textctrl.Bind(wx.EVT_TEXT_ENTER, self.xsize_textctrl_entered)
self.x1_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.x1_textctrl.SetFont(self.textentry_font)
values_sizer.Add(self.x1_textctrl, 0, wx.ALL, 2)
self.x1_textctrl.Bind(wx.EVT_TEXT, self.x1_textctrl_changed)
self.x1_textctrl.Bind(wx.EVT_TEXT_ENTER, self.x1_textctrl_entered)
self.npix_static_text = wx.StaticText( self, wx.ID_ANY, u"# pixels", wx.DefaultPosition, wx.DefaultSize, 0 )
self.npix_static_text.Wrap( -1 )
values_sizer.Add(self.npix_static_text, 0, wx.ALL|wx.ALIGN_CENTER_HORIZONTAL|wx.ALIGN_BOTTOM, 0)
self.y_static_text = wx.StaticText( self, wx.ID_ANY, u"y", wx.DefaultPosition, wx.DefaultSize, 0 )
self.y_static_text.Wrap( -1 )
values_sizer.Add(self.y_static_text, 0, wx.ALL|wx.ALIGN_CENTER_VERTICAL, 0)
self.y0_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.y0_textctrl.SetFont(self.textentry_font)
values_sizer.Add(self.y0_textctrl, 0, wx.ALL, 2)
self.y0_textctrl.Bind(wx.EVT_TEXT, self.y0_textctrl_changed)
self.y0_textctrl.Bind(wx.EVT_TEXT_ENTER, self.y0_textctrl_entered)
self.ysize_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.ysize_textctrl.SetFont(self.textentry_font)
values_sizer.Add(self.ysize_textctrl, 0, wx.ALL, 2)
self.ysize_textctrl.Bind(wx.EVT_TEXT, self.ysize_textctrl_changed)
self.ysize_textctrl.Bind(wx.EVT_TEXT_ENTER, self.ysize_textctrl_entered)
self.y1_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.y1_textctrl.SetFont(self.textentry_font)
values_sizer.Add(self.y1_textctrl, 0, wx.ALL, 2)
self.y1_textctrl.Bind(wx.EVT_TEXT, self.y1_textctrl_changed)
self.y1_textctrl.Bind(wx.EVT_TEXT_ENTER, self.y1_textctrl_entered)
self.npix_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.npix_textctrl.SetFont(self.textentry_font)
self.npix_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.npix_textctrl, 0, wx.ALL|wx.ALIGN_CENTER_VERTICAL|wx.ALIGN_LEFT, 0)
values_sizer.AddSpacer((0,15), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
self.median_static_text = wx.StaticText( self, wx.ID_ANY, u"Median", wx.DefaultPosition, wx.DefaultSize, 0 )
self.median_static_text.Wrap( -1 )
values_sizer.Add(self.median_static_text, 0, wx.ALL|wx.ALIGN_CENTER_VERTICAL|wx.ALIGN_RIGHT, 0)
self.median_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.median_textctrl.SetFont(self.textentry_font)
self.median_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.median_textctrl, 0, wx.ALL, 2)
self.robust_static_text = wx.StaticText( self, wx.ID_ANY, u"Robust", wx.DefaultPosition, wx.DefaultSize, 0 )
self.robust_static_text.Wrap( -1 )
values_sizer.Add(self.robust_static_text, 0, wx.ALL|wx.ALIGN_BOTTOM|wx.ALIGN_CENTER_HORIZONTAL, 0)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
self.mean_static_text = wx.StaticText( self, wx.ID_ANY, u"Mean", wx.DefaultPosition, wx.DefaultSize, 0 )
self.mean_static_text.Wrap( -1 )
values_sizer.Add(self.mean_static_text, 0, wx.ALL|wx.ALIGN_CENTER_VERTICAL|wx.ALIGN_RIGHT, 0)
self.mean_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.mean_textctrl.SetFont(self.textentry_font)
self.mean_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.mean_textctrl, 0, wx.ALL, 2)
self.robust_mean_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.robust_mean_textctrl.SetFont(self.textentry_font)
self.robust_mean_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.robust_mean_textctrl, 0, wx.ALL, 2)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
self.stdev_static_text = wx.StaticText( self, wx.ID_ANY, u"Stdev", wx.DefaultPosition, wx.DefaultSize, 0 )
self.stdev_static_text.Wrap( -1 )
values_sizer.Add(self.stdev_static_text, 0, wx.ALL|wx.ALIGN_CENTER_VERTICAL|wx.ALIGN_RIGHT, 0)
self.stdev_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.stdev_textctrl.SetFont(self.textentry_font)
self.stdev_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.stdev_textctrl, 0, wx.ALL, 2)
self.robust_stdev_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.robust_stdev_textctrl.SetFont(self.textentry_font)
self.robust_stdev_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.robust_stdev_textctrl, 0, wx.ALL, 2)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,15), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
self.min_static_text = wx.StaticText( self, wx.ID_ANY, u"Min", wx.DefaultPosition, wx.DefaultSize, 0 )
self.min_static_text.Wrap( -1 )
values_sizer.Add(self.min_static_text, 0, wx.ALL|wx.ALIGN_CENTER_VERTICAL|wx.ALIGN_RIGHT, 0)
self.minval_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.minval_textctrl.SetFont(self.textentry_font)
self.minval_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.minval_textctrl, 0, wx.ALL, 2)
self.minpos_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.minpos_textctrl.SetFont(self.textentry_font)
self.minpos_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.minpos_textctrl, 0, wx.ALL, 2)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
values_sizer.AddSpacer((0,0), 0, wx.EXPAND)
self.max_static_text = wx.StaticText( self, wx.ID_ANY, u"Max", wx.DefaultPosition, wx.DefaultSize, 0 )
self.max_static_text.Wrap( -1 )
values_sizer.Add(self.max_static_text, 0, wx.ALL|wx.ALIGN_CENTER_VERTICAL|wx.ALIGN_RIGHT, 0)
self.maxval_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.maxval_textctrl.SetFont(self.textentry_font)
self.maxval_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.maxval_textctrl, 0, wx.ALL, 2)
self.maxpos_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.maxpos_textctrl.SetFont(self.textentry_font)
self.maxpos_textctrl.SetBackgroundColour(textctrl_output_only_background_color)
values_sizer.Add(self.maxpos_textctrl, 0, wx.ALL, 2)
self.hideshow_button = wx.Button(self, wx.ID_ANY, u"Show", wx.DefaultPosition, wx.DefaultSize, 0)
values_sizer.Add(self.hideshow_button, 0, wx.ALL|wx.ALIGN_CENTER_HORIZONTAL|wx.ALIGN_CENTER_VERTICAL, 2)
self.hideshow_button.Bind(wx.EVT_BUTTON, self.on_hideshow_button)
v_sizer1 = wx.BoxSizer(wx.VERTICAL)
v_sizer1.AddStretchSpacer(1.0)
v_sizer1.Add(values_sizer, 0, wx.ALIGN_CENTER_HORIZONTAL)
v_sizer1.AddStretchSpacer(1.0)
self.SetSizer(v_sizer1)
pub.subscribe(self.queue_update_stats, 'recalc-display-image-called')
pub.subscribe(self._set_stats_box_parameters, 'set-stats-box-parameters')
pub.subscribe(self.publish_stats_to_stream, 'get-stats-box-info')
def publish_stats_to_stream(self, msg=None):
wx.CallAfter(send_to_stream, sys.stdout, ('stats-box-info', self.stats_info))
def on_button_press(self, event):
self.select_panel()
self.update_stats_box(event.xdata, event.ydata, event.xdata, event.ydata)
self.redraw_overplot_on_image()
self.cursor_stats_box_x0, self.cursor_stats_box_y0 = event.xdata, event.ydata
def on_motion(self, event):
if event.button is not None:
self.update_stats_box(self.cursor_stats_box_x0, self.cursor_stats_box_y0, event.xdata, event.ydata)
self.redraw_overplot_on_image()
self.update_stats()
def on_button_release(self, event):
self.redraw_overplot_on_image()
self.update_stats()
def set_cursor_to_stats_box_mode(self, event):
self.ztv_frame.primary_image_panel.cursor_mode = 'Stats box'
self.ztv_frame.stats_panel.select_panel()
self.ztv_frame.stats_panel.highlight_panel()
def queue_update_stats(self, msg=None):
"""
wrapper to call update_stats from CallAfter in order to make GUI as responsive as possible.
"""
wx.CallAfter(self.update_stats, msg=None)
def _set_stats_box_parameters(self, msg):
"""
wrapper to update_stats_box to receive messages & translate them correctly
"""
x0,x1,y0,y1 = [None]*4
if msg['xrange'] is not None:
x0,x1 = msg['xrange']
if msg['yrange'] is not None:
y0,y1 = msg['yrange']
if msg['xrange'] is not None or msg['yrange'] is not None:
self.update_stats_box(x0, y0, x1, y1)
if msg['show_overplot'] is not None:
if msg['show_overplot']:
self.redraw_overplot_on_image()
else:
self.remove_overplot_on_image()
send_to_stream(sys.stdout, ('set-stats-box-parameters-done', True))
def update_stats_box(self, x0=None, y0=None, x1=None, y1=None):
if x0 is None:
x0 = self.stats_rect.get_x()
if y0 is None:
y0 = self.stats_rect.get_y()
if x1 is None:
x1 = self.stats_rect.get_x() + self.stats_rect.get_width()
if y1 is None:
y1 = self.stats_rect.get_y() + self.stats_rect.get_height()
if x0 > x1:
x0, x1 = x1, x0
if y0 > y1:
y0, y1 = y1, y0
x0 = min(max(0, x0), self.ztv_frame.display_image.shape[1] - 1)
y0 = min(max(0, y0), self.ztv_frame.display_image.shape[0] - 1)
x1 = min(max(0, x1), self.ztv_frame.display_image.shape[1] - 1)
y1 = min(max(0, y1), self.ztv_frame.display_image.shape[0] - 1)
self.stats_rect.set_bounds(x0, y0, x1 - x0, y1 - y0)
if self.hideshow_button.GetLabel() == 'Hide':
self.ztv_frame.primary_image_panel.figure.canvas.draw()
self.update_stats()
def remove_overplot_on_image(self):
self.ztv_frame.primary_image_panel.remove_patch('stats_panel:stats_rect')
self.hideshow_button.SetLabel(u"Show")
def redraw_overplot_on_image(self):
self.ztv_frame.primary_image_panel.add_patch('stats_panel:stats_rect', self.stats_rect)
self.hideshow_button.SetLabel(u"Hide")
def on_hideshow_button(self, evt):
if self.hideshow_button.GetLabel() == 'Hide':
self.remove_overplot_on_image()
else:
self.redraw_overplot_on_image()
def get_x0y0x1y1_from_stats_rect(self):
x0 = self.stats_rect.get_x()
y0 = self.stats_rect.get_y()
x1 = x0 + self.stats_rect.get_width()
y1 = y0 + self.stats_rect.get_height()
return x0,y0,x1,y1
def update_stats(self, msg=None):
x0,y0,x1,y1 = self.get_x0y0x1y1_from_stats_rect()
x0, y0 = int(np.round(x0)), int(np.round(y0))
x1, y1 = int(np.round(x1)), int(np.round(y1))
self.last_string_values['x0'] = str(int(x0))
self.x0_textctrl.SetValue(self.last_string_values['x0'])
self.last_string_values['y0'] = str(int(y0))
self.y0_textctrl.SetValue(self.last_string_values['y0'])
x_npix = int(x1 - x0 + 1)
self.last_string_values['xsize'] = str(x_npix)
self.xsize_textctrl.SetValue(self.last_string_values['xsize'])
y_npix = int(y1 - y0 + 1)
self.last_string_values['ysize'] = str(y_npix)
self.ysize_textctrl.SetValue(self.last_string_values['ysize'])
self.last_string_values['x1'] = str(int(x1))
self.x1_textctrl.SetValue(self.last_string_values['x1'])
self.last_string_values['y1'] = str(int(y1))
self.y1_textctrl.SetValue(self.last_string_values['y1'])
self.npix_textctrl.SetValue(str(x_npix * y_npix))
stats_data = self.ztv_frame.display_image[y0:y1+1, x0:x1+1]
finite_mask = np.isfinite(stats_data)
if finite_mask.max() is np.True_:
stats_data_mean = stats_data[finite_mask].mean()
stats_data_median = np.median(stats_data[finite_mask])
stats_data_std = stats_data[finite_mask].std()
robust_mean, robust_median, robust_std = sigma_clipped_stats(stats_data[finite_mask])
else:
stats_data_mean = np.nan
stats_data_median = np.nan
stats_data_std = np.inf
robust_mean, robust_median, robust_std = np.nan, np.nan, np.inf
self.stats_info = {'xrange':[x0,x1], 'yrange':[y0,y1],
'mean':stats_data_mean, 'median':stats_data_median, 'std':stats_data_std,
'min':stats_data.min(), 'max':stats_data.max()} # want min/max to reflect any Inf/NaN
self.mean_textctrl.SetValue("{:0.4g}".format(self.stats_info['mean']))
self.median_textctrl.SetValue("{:0.4g}".format(self.stats_info['median']))
self.stdev_textctrl.SetValue("{:0.4g}".format(self.stats_info['std']))
self.stats_info['robust-mean'] = robust_mean
self.stats_info['robust-median'] = robust_median
self.stats_info['robust-std'] = robust_std
self.robust_mean_textctrl.SetValue("{:0.4g}".format(robust_mean))
self.robust_stdev_textctrl.SetValue("{:0.4g}".format(robust_std))
self.minval_textctrl.SetValue("{:0.4g}".format(self.stats_info['min']))
self.maxval_textctrl.SetValue("{:0.4g}".format(self.stats_info['max']))
wmin = np.where(stats_data == stats_data.min())
wmin = [(wmin[1][i] + x0,wmin[0][i] + y0) for i in np.arange(wmin[0].size)]
if len(wmin) == 1:
wmin = wmin[0]
self.minpos_textctrl.SetValue("{}".format(wmin))
self.stats_info['wmin'] = wmin
wmax = np.where(stats_data == stats_data.max())
wmax = [(wmax[1][i] + x0,wmax[0][i] + y0) for i in np.arange(wmax[0].size)]
if len(wmax) == 1:
wmax = wmax[0]
self.maxpos_textctrl.SetValue("{}".format(wmax))
self.stats_info['wmax'] = wmax
set_textctrl_background_color(self.x0_textctrl, 'ok')
set_textctrl_background_color(self.x1_textctrl, 'ok')
set_textctrl_background_color(self.xsize_textctrl, 'ok')
set_textctrl_background_color(self.y0_textctrl, 'ok')
set_textctrl_background_color(self.y1_textctrl, 'ok')
set_textctrl_background_color(self.ysize_textctrl, 'ok')
def x0_textctrl_changed(self, evt):
validate_textctrl_str(self.x0_textctrl, int, self.last_string_values['x0'])
def x0_textctrl_entered(self, evt):
if validate_textctrl_str(self.x0_textctrl, int, self.last_string_values['x0']):
self.last_string_values['x0'] = self.x0_textctrl.GetValue()
self.update_stats_box(int(self.last_string_values['x0']), None, None, None)
self.x0_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def xsize_textctrl_changed(self, evt):
validate_textctrl_str(self.xsize_textctrl, int, self.last_string_values['xsize'])
def xsize_textctrl_entered(self, evt):
if validate_textctrl_str(self.xsize_textctrl, int, self.last_string_values['xsize']):
self.last_string_values['xsize'] = self.xsize_textctrl.GetValue()
xsize = int(self.last_string_values['xsize'])
sys.stderr.write("\n\nxsize = {}\n\n".format(xsize))
x0,y0,x1,y1 = self.get_x0y0x1y1_from_stats_rect()
xc = (x0 + x1) / 2.
x0 = max(0, int(xc - xsize / 2.))
x1 = x0 + xsize - 1
x1 = min(x1, self.ztv_frame.display_image.shape[1] - 1)
x0 = x1 - xsize + 1
x0 = max(0, int(xc - xsize / 2.))
self.update_stats_box(x0, y0, x1, y1)
self.xsize_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def x1_textctrl_changed(self, evt):
validate_textctrl_str(self.x1_textctrl, int, self.last_string_values['x1'])
def x1_textctrl_entered(self, evt):
if validate_textctrl_str(self.x1_textctrl, int, self.last_string_values['x1']):
self.last_string_values['x1'] = self.x1_textctrl.GetValue()
self.update_stats_box(None, None, int(self.last_string_values['x1']), None)
self.x1_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def y0_textctrl_changed(self, evt):
validate_textctrl_str(self.y0_textctrl, int, self.last_string_values['y0'])
def y0_textctrl_entered(self, evt):
if validate_textctrl_str(self.y0_textctrl, int, self.last_string_values['y0']):
self.last_string_values['y0'] = self.y0_textctrl.GetValue()
self.update_stats_box(None, int(self.last_string_values['y0']), None, None)
self.y0_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def ysize_textctrl_changed(self, evt):
validate_textctrl_str(self.ysize_textctrl, int, self.last_string_values['ysize'])
def ysize_textctrl_entered(self, evt):
if validate_textctrl_str(self.ysize_textctrl, int, self.last_string_values['ysize']):
self.last_string_values['ysize'] = self.ysize_textctrl.GetValue()
ysize = int(self.last_string_values['ysize'])
x0,y0,x1,y1 = self.get_x0y0x1y1_from_stats_rect()
yc = (y0 + y1) / 2.
y0 = max(0, int(yc - ysize / 2.))
y1 = y0 + ysize - 1
y1 = min(y1, self.ztv_frame.display_image.shape[0] - 1)
y0 = y1 - ysize + 1
y0 = max(0, int(yc - ysize / 2.))
self.update_stats_box(x0, y0, x1, y1)
self.ysize_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def y1_textctrl_changed(self, evt):
validate_textctrl_str(self.y1_textctrl, int, self.last_string_values['y1'])
def y1_textctrl_entered(self, evt):
if validate_textctrl_str(self.y1_textctrl, int, self.last_string_values['y1']):
self.last_string_values['y1'] = self.y1_textctrl.GetValue()
self.update_stats_box(None, None, None, int(self.last_string_values['y1']))
self.y1_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
class StatsPanel(wx.Panel):
def __init__(self, parent):
wx.Panel.__init__(self, parent, wx.ID_ANY, wx.DefaultPosition,
wx.DefaultSize)
self.ztv_frame = self.GetTopLevelParent()
self.ztv_frame.primary_image_panel.popup_menu_cursor_modes.append(
'Stats box')
self.ztv_frame.primary_image_panel.available_cursor_modes[
'Stats box'] = {
'set-to-mode': self.set_cursor_to_stats_box_mode,
'on_button_press': self.on_button_press,
'on_motion': self.on_motion,
'on_button_release': self.on_button_release
}
self.stats_info = None
self.last_string_values = {
'x0': '',
'xsize': '',
'x1': '',
'y0': '',
'ysize': '',
'y1': ''
}
self.stats_rect = Rectangle((0, 0),
10,
10,
color='magenta',
fill=False,
zorder=100)
# use self.stats_rect as where we store/retrieve the x0,y0,x1,y1
# x0,y0,x1,y1 should be limited to range of 0 to shape-1
# but, stats should be calculated over e.g. x0:x1+1 (so that have pixels to do stats on even if x0==x1)
# and, width/height of stats_rect should always be >= 0
textentry_font = wx.Font(14, wx.FONTFAMILY_MODERN, wx.NORMAL,
wx.FONTWEIGHT_LIGHT, False)
values_sizer = wx.FlexGridSizer(10, 5, 0, 0)
values_sizer.SetFlexibleDirection(wx.BOTH)
values_sizer.SetNonFlexibleGrowMode(wx.FLEX_GROWMODE_SPECIFIED)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
self.low_static_text = wx.StaticText(self, wx.ID_ANY, u"Low",
wx.DefaultPosition,
wx.DefaultSize, wx.ALIGN_RIGHT)
self.low_static_text.Wrap(-1)
values_sizer.Add(self.low_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_HORIZONTAL, 0)
self.low_static_text = wx.StaticText(self, wx.ID_ANY, u"# pix",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.low_static_text.Wrap(-1)
values_sizer.Add(self.low_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_HORIZONTAL, 0)
self.high_static_text = wx.StaticText(self, wx.ID_ANY, u"High",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.high_static_text.Wrap(-1)
values_sizer.Add(self.high_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_HORIZONTAL, 0)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
self.x_static_text = wx.StaticText(self, wx.ID_ANY, u"x",
wx.DefaultPosition, wx.DefaultSize,
0)
self.x_static_text.Wrap(-1)
values_sizer.Add(self.x_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_VERTICAL, 0)
self.x0_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.x0_textctrl.SetFont(textentry_font)
values_sizer.Add(self.x0_textctrl, 0, wx.ALL, 2)
self.x0_textctrl.Bind(wx.EVT_TEXT, self.x0_textctrl_changed)
self.x0_textctrl.Bind(wx.EVT_TEXT_ENTER, self.x0_textctrl_entered)
self.xsize_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.xsize_textctrl.SetFont(textentry_font)
values_sizer.Add(self.xsize_textctrl, 0, wx.ALL, 2)
self.xsize_textctrl.Bind(wx.EVT_TEXT, self.xsize_textctrl_changed)
self.xsize_textctrl.Bind(wx.EVT_TEXT_ENTER,
self.xsize_textctrl_entered)
self.x1_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.x1_textctrl.SetFont(textentry_font)
values_sizer.Add(self.x1_textctrl, 0, wx.ALL, 2)
self.x1_textctrl.Bind(wx.EVT_TEXT, self.x1_textctrl_changed)
self.x1_textctrl.Bind(wx.EVT_TEXT_ENTER, self.x1_textctrl_entered)
self.npix_static_text = wx.StaticText(self, wx.ID_ANY, u"# pixels",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.npix_static_text.Wrap(-1)
values_sizer.Add(self.npix_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_HORIZONTAL | wx.ALIGN_BOTTOM,
0)
self.y_static_text = wx.StaticText(self, wx.ID_ANY, u"y",
wx.DefaultPosition, wx.DefaultSize,
0)
self.y_static_text.Wrap(-1)
values_sizer.Add(self.y_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_VERTICAL, 0)
self.y0_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.y0_textctrl.SetFont(textentry_font)
values_sizer.Add(self.y0_textctrl, 0, wx.ALL, 2)
self.y0_textctrl.Bind(wx.EVT_TEXT, self.y0_textctrl_changed)
self.y0_textctrl.Bind(wx.EVT_TEXT_ENTER, self.y0_textctrl_entered)
self.ysize_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.ysize_textctrl.SetFont(textentry_font)
values_sizer.Add(self.ysize_textctrl, 0, wx.ALL, 2)
self.ysize_textctrl.Bind(wx.EVT_TEXT, self.ysize_textctrl_changed)
self.ysize_textctrl.Bind(wx.EVT_TEXT_ENTER,
self.ysize_textctrl_entered)
self.y1_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_PROCESS_ENTER)
self.y1_textctrl.SetFont(textentry_font)
values_sizer.Add(self.y1_textctrl, 0, wx.ALL, 2)
self.y1_textctrl.Bind(wx.EVT_TEXT, self.y1_textctrl_changed)
self.y1_textctrl.Bind(wx.EVT_TEXT_ENTER, self.y1_textctrl_entered)
self.npix_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.npix_textctrl.SetFont(textentry_font)
self.npix_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.npix_textctrl, 0,
wx.ALL | wx.ALIGN_CENTER_VERTICAL | wx.ALIGN_LEFT, 0)
values_sizer.AddSpacer((0, 15), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
self.median_static_text = wx.StaticText(self, wx.ID_ANY, u"Median",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.median_static_text.Wrap(-1)
values_sizer.Add(self.median_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_VERTICAL | wx.ALIGN_RIGHT, 0)
self.median_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.median_textctrl.SetFont(textentry_font)
self.median_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.median_textctrl, 0, wx.ALL, 2)
self.robust_static_text = wx.StaticText(self, wx.ID_ANY, u"Robust",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.robust_static_text.Wrap(-1)
values_sizer.Add(self.robust_static_text, 0,
wx.ALL | wx.ALIGN_BOTTOM | wx.ALIGN_CENTER_HORIZONTAL,
0)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
self.mean_static_text = wx.StaticText(self, wx.ID_ANY, u"Mean",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.mean_static_text.Wrap(-1)
values_sizer.Add(self.mean_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_VERTICAL | wx.ALIGN_RIGHT, 0)
self.mean_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.mean_textctrl.SetFont(textentry_font)
self.mean_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.mean_textctrl, 0, wx.ALL, 2)
self.robust_mean_textctrl = wx.TextCtrl(self, wx.ID_ANY,
wx.EmptyString,
wx.DefaultPosition,
wx.DefaultSize, wx.TE_READONLY)
self.robust_mean_textctrl.SetFont(textentry_font)
self.robust_mean_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.robust_mean_textctrl, 0, wx.ALL, 2)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
self.stdev_static_text = wx.StaticText(self, wx.ID_ANY, u"Stdev",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.stdev_static_text.Wrap(-1)
values_sizer.Add(self.stdev_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_VERTICAL | wx.ALIGN_RIGHT, 0)
self.stdev_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.stdev_textctrl.SetFont(textentry_font)
self.stdev_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.stdev_textctrl, 0, wx.ALL, 2)
self.robust_stdev_textctrl = wx.TextCtrl(self, wx.ID_ANY,
wx.EmptyString,
wx.DefaultPosition,
wx.DefaultSize,
wx.TE_READONLY)
self.robust_stdev_textctrl.SetFont(textentry_font)
self.robust_stdev_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.robust_stdev_textctrl, 0, wx.ALL, 2)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 15), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
self.min_static_text = wx.StaticText(self, wx.ID_ANY, u"Min",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.min_static_text.Wrap(-1)
values_sizer.Add(self.min_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_VERTICAL | wx.ALIGN_RIGHT, 0)
self.minval_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.minval_textctrl.SetFont(textentry_font)
self.minval_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.minval_textctrl, 0, wx.ALL, 2)
self.minpos_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.minpos_textctrl.SetFont(textentry_font)
self.minpos_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.minpos_textctrl, 0, wx.ALL, 2)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
values_sizer.AddSpacer((0, 0), 0, wx.EXPAND)
self.max_static_text = wx.StaticText(self, wx.ID_ANY, u"Max",
wx.DefaultPosition,
wx.DefaultSize, 0)
self.max_static_text.Wrap(-1)
values_sizer.Add(self.max_static_text, 0,
wx.ALL | wx.ALIGN_CENTER_VERTICAL | wx.ALIGN_RIGHT, 0)
self.maxval_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.maxval_textctrl.SetFont(textentry_font)
self.maxval_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.maxval_textctrl, 0, wx.ALL, 2)
self.maxpos_textctrl = wx.TextCtrl(self, wx.ID_ANY, wx.EmptyString,
wx.DefaultPosition, wx.DefaultSize,
wx.TE_READONLY)
self.maxpos_textctrl.SetFont(textentry_font)
self.maxpos_textctrl.SetBackgroundColour(
textctrl_output_only_background_color)
values_sizer.Add(self.maxpos_textctrl, 0, wx.ALL, 2)
self.hideshow_button = wx.Button(self, wx.ID_ANY, u"Show",
wx.DefaultPosition, wx.DefaultSize, 0)
values_sizer.Add(
self.hideshow_button, 0,
wx.ALL | wx.ALIGN_CENTER_HORIZONTAL | wx.ALIGN_CENTER_VERTICAL, 2)
self.hideshow_button.Bind(wx.EVT_BUTTON, self.on_hideshow_button)
v_sizer1 = wx.BoxSizer(wx.VERTICAL)
v_sizer1.AddStretchSpacer(1.0)
v_sizer1.Add(values_sizer, 0, wx.ALIGN_CENTER_HORIZONTAL)
v_sizer1.AddStretchSpacer(1.0)
self.SetSizer(v_sizer1)
pub.subscribe(self.queue_update_stats, 'recalc-display-image-called')
pub.subscribe(self._set_stats_box_parameters,
'set-stats-box-parameters')
pub.subscribe(self.publish_stats_to_stream, 'get-stats-box-info')
def publish_stats_to_stream(self, msg=None):
wx.CallAfter(send_to_stream, sys.stdout,
('stats-box-info', self.stats_info))
def on_button_press(self, event):
self.select_panel()
self.stats_start_timestamp = event.guiEvent.GetTimestamp() # millisec
self.update_stats_box(event.xdata, event.ydata, event.xdata,
event.ydata)
self.redraw_overplot_on_image()
self.cursor_stats_box_x0, self.cursor_stats_box_y0 = event.xdata, event.ydata
def on_motion(self, event):
self.update_stats_box(self.cursor_stats_box_x0,
self.cursor_stats_box_y0, event.xdata,
event.ydata)
self.redraw_overplot_on_image()
self.update_stats()
def on_button_release(self, event):
self.redraw_overplot_on_image()
self.update_stats()
def set_cursor_to_stats_box_mode(self, event):
self.ztv_frame.primary_image_panel.cursor_mode = 'Stats box'
self.ztv_frame.stats_panel.select_panel()
self.ztv_frame.stats_panel.highlight_panel()
def queue_update_stats(self, msg=None):
"""
wrapper to call update_stats from CallAfter in order to make GUI as responsive as possible.
"""
wx.CallAfter(self.update_stats, msg=None)
def _set_stats_box_parameters(self, msg):
"""
wrapper to update_stats_box to receive messages & translate them correctly
"""
x0, x1, y0, y1 = [None] * 4
if msg['xrange'] is not None:
x0, x1 = msg['xrange']
if msg['yrange'] is not None:
y0, y1 = msg['yrange']
if msg['xrange'] is not None or msg['yrange'] is not None:
self.update_stats_box(x0, y0, x1, y1)
if msg['show_overplot'] is not None:
if msg['show_overplot']:
self.redraw_overplot_on_image()
else:
self.remove_overplot_on_image()
send_to_stream(sys.stdout, ('set-stats-box-parameters-done', True))
def update_stats_box(self, x0=None, y0=None, x1=None, y1=None):
if x0 is None:
x0 = self.stats_rect.get_x()
if y0 is None:
y0 = self.stats_rect.get_y()
if x1 is None:
x1 = self.stats_rect.get_x() + self.stats_rect.get_width()
if y1 is None:
y1 = self.stats_rect.get_y() + self.stats_rect.get_height()
if x0 > x1:
x0, x1 = x1, x0
if y0 > y1:
y0, y1 = y1, y0
x0 = min(max(0, x0), self.ztv_frame.display_image.shape[1] - 1)
y0 = min(max(0, y0), self.ztv_frame.display_image.shape[0] - 1)
x1 = min(max(0, x1), self.ztv_frame.display_image.shape[1] - 1)
y1 = min(max(0, y1), self.ztv_frame.display_image.shape[0] - 1)
self.stats_rect.set_bounds(x0, y0, x1 - x0, y1 - y0)
self.ztv_frame.primary_image_panel.figure.canvas.draw()
self.update_stats()
def remove_overplot_on_image(self):
if self.stats_rect in self.ztv_frame.primary_image_panel.axes.patches:
self.ztv_frame.primary_image_panel.axes.patches.remove(
self.stats_rect)
self.ztv_frame.primary_image_panel.figure.canvas.draw()
self.hideshow_button.SetLabel(u"Show")
def redraw_overplot_on_image(self):
if self.stats_rect not in self.ztv_frame.primary_image_panel.axes.patches:
self.ztv_frame.primary_image_panel.axes.add_patch(self.stats_rect)
self.ztv_frame.primary_image_panel.figure.canvas.draw()
self.hideshow_button.SetLabel(u"Hide")
def on_hideshow_button(self, evt):
if self.hideshow_button.GetLabel() == 'Hide':
self.remove_overplot_on_image()
else:
self.redraw_overplot_on_image()
def get_x0y0x1y1_from_stats_rect(self):
x0 = self.stats_rect.get_x()
y0 = self.stats_rect.get_y()
x1 = x0 + self.stats_rect.get_width()
y1 = y0 + self.stats_rect.get_height()
return x0, y0, x1, y1
def update_stats(self, msg=None):
x0, y0, x1, y1 = self.get_x0y0x1y1_from_stats_rect()
x0, y0 = int(np.round(x0)), int(np.round(y0))
x1, y1 = int(np.round(x1)), int(np.round(y1))
self.last_string_values['x0'] = str(int(x0))
self.x0_textctrl.SetValue(self.last_string_values['x0'])
self.last_string_values['y0'] = str(int(y0))
self.y0_textctrl.SetValue(self.last_string_values['y0'])
x_npix = int(x1 - x0 + 1)
self.last_string_values['xsize'] = str(x_npix)
self.xsize_textctrl.SetValue(self.last_string_values['xsize'])
y_npix = int(y1 - y0 + 1)
self.last_string_values['ysize'] = str(y_npix)
self.ysize_textctrl.SetValue(self.last_string_values['ysize'])
self.last_string_values['x1'] = str(int(x1))
self.x1_textctrl.SetValue(self.last_string_values['x1'])
self.last_string_values['y1'] = str(int(y1))
self.y1_textctrl.SetValue(self.last_string_values['y1'])
self.npix_textctrl.SetValue(str(x_npix * y_npix))
stats_data = self.ztv_frame.display_image[y0:y1 + 1, x0:x1 + 1]
finite_mask = np.isfinite(stats_data)
if finite_mask.max() is np.True_:
stats_data_mean = stats_data[finite_mask].mean()
stats_data_median = np.median(stats_data[finite_mask])
stats_data_std = stats_data[finite_mask].std()
robust_mean, robust_median, robust_std = sigma_clipped_stats(
stats_data[finite_mask])
else:
stats_data_mean = np.nan
stats_data_median = np.nan
stats_data_std = np.inf
robust_mean, robust_median, robust_std = np.nan, np.nan, np.inf
self.stats_info = {
'xrange': [x0, x1],
'yrange': [y0, y1],
'mean': stats_data_mean,
'median': stats_data_median,
'std': stats_data_std,
'min': stats_data.min(),
'max': stats_data.max()
} # want min/max to reflect any Inf/NaN
self.mean_textctrl.SetValue("{:0.4g}".format(self.stats_info['mean']))
self.median_textctrl.SetValue("{:0.4g}".format(
self.stats_info['median']))
self.stdev_textctrl.SetValue("{:0.4g}".format(self.stats_info['std']))
self.stats_info['robust-mean'] = robust_mean
self.stats_info['robust-median'] = robust_median
self.stats_info['robust-std'] = robust_std
self.robust_mean_textctrl.SetValue("{:0.4g}".format(robust_mean))
self.robust_stdev_textctrl.SetValue("{:0.4g}".format(robust_std))
self.minval_textctrl.SetValue("{:0.4g}".format(self.stats_info['min']))
self.maxval_textctrl.SetValue("{:0.4g}".format(self.stats_info['max']))
wmin = np.where(stats_data == stats_data.min())
wmin = [(wmin[1][i] + x0, wmin[0][i] + y0)
for i in np.arange(wmin[0].size)]
if len(wmin) == 1:
wmin = wmin[0]
self.minpos_textctrl.SetValue("{}".format(wmin))
self.stats_info['wmin'] = wmin
wmax = np.where(stats_data == stats_data.max())
wmax = [(wmax[1][i] + x0, wmax[0][i] + y0)
for i in np.arange(wmax[0].size)]
if len(wmax) == 1:
wmax = wmax[0]
self.maxpos_textctrl.SetValue("{}".format(wmax))
self.stats_info['wmax'] = wmax
set_textctrl_background_color(self.x0_textctrl, 'ok')
set_textctrl_background_color(self.x1_textctrl, 'ok')
set_textctrl_background_color(self.xsize_textctrl, 'ok')
set_textctrl_background_color(self.y0_textctrl, 'ok')
set_textctrl_background_color(self.y1_textctrl, 'ok')
set_textctrl_background_color(self.ysize_textctrl, 'ok')
def x0_textctrl_changed(self, evt):
validate_textctrl_str(self.x0_textctrl, int,
self.last_string_values['x0'])
def x0_textctrl_entered(self, evt):
if validate_textctrl_str(self.x0_textctrl, int,
self.last_string_values['x0']):
self.last_string_values['x0'] = self.x0_textctrl.GetValue()
self.update_stats_box(int(self.last_string_values['x0']), None,
None, None)
self.x0_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def xsize_textctrl_changed(self, evt):
validate_textctrl_str(self.xsize_textctrl, int,
self.last_string_values['xsize'])
def xsize_textctrl_entered(self, evt):
if validate_textctrl_str(self.xsize_textctrl, int,
self.last_string_values['xsize']):
self.last_string_values['xsize'] = self.xsize_textctrl.GetValue()
xsize = int(self.last_string_values['xsize'])
sys.stderr.write("\n\nxsize = {}\n\n".format(xsize))
x0, y0, x1, y1 = self.get_x0y0x1y1_from_stats_rect()
xc = (x0 + x1) / 2.
x0 = max(0, int(xc - xsize / 2.))
x1 = x0 + xsize - 1
x1 = min(x1, self.ztv_frame.display_image.shape[1] - 1)
x0 = x1 - xsize + 1
x0 = max(0, int(xc - xsize / 2.))
self.update_stats_box(x0, y0, x1, y1)
self.xsize_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def x1_textctrl_changed(self, evt):
validate_textctrl_str(self.x1_textctrl, int,
self.last_string_values['x1'])
def x1_textctrl_entered(self, evt):
if validate_textctrl_str(self.x1_textctrl, int,
self.last_string_values['x1']):
self.last_string_values['x1'] = self.x1_textctrl.GetValue()
self.update_stats_box(None, None,
int(self.last_string_values['x1']), None)
self.x1_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def y0_textctrl_changed(self, evt):
validate_textctrl_str(self.y0_textctrl, int,
self.last_string_values['y0'])
def y0_textctrl_entered(self, evt):
if validate_textctrl_str(self.y0_textctrl, int,
self.last_string_values['y0']):
self.last_string_values['y0'] = self.y0_textctrl.GetValue()
self.update_stats_box(None, int(self.last_string_values['y0']),
None, None)
self.y0_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def ysize_textctrl_changed(self, evt):
validate_textctrl_str(self.ysize_textctrl, int,
self.last_string_values['ysize'])
def ysize_textctrl_entered(self, evt):
if validate_textctrl_str(self.ysize_textctrl, int,
self.last_string_values['ysize']):
self.last_string_values['ysize'] = self.ysize_textctrl.GetValue()
ysize = int(self.last_string_values['ysize'])
x0, y0, x1, y1 = self.get_x0y0x1y1_from_stats_rect()
yc = (y0 + y1) / 2.
y0 = max(0, int(yc - ysize / 2.))
y1 = y0 + ysize - 1
y1 = min(y1, self.ztv_frame.display_image.shape[0] - 1)
y0 = y1 - ysize + 1
y0 = max(0, int(yc - ysize / 2.))
self.update_stats_box(x0, y0, x1, y1)
self.ysize_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
def y1_textctrl_changed(self, evt):
validate_textctrl_str(self.y1_textctrl, int,
self.last_string_values['y1'])
def y1_textctrl_entered(self, evt):
if validate_textctrl_str(self.y1_textctrl, int,
self.last_string_values['y1']):
self.last_string_values['y1'] = self.y1_textctrl.GetValue()
self.update_stats_box(None, None, None,
int(self.last_string_values['y1']))
self.y1_textctrl.SetSelection(-1, -1)
self.redraw_overplot_on_image()
class SimGui(wx.Frame):
StateNone, StateSetStartPos, StateSetGoalPos = range(3)
SimStateStopped, SimStateRunning, SimStatePaused = range(3)
def __init__(self, size=(1800, 1200)):
super(SimGui, self).__init__(None, size=size)
self.plotter = Plot(self)
self.map = None
self.map_visualizer = None
self.Bind(wx.EVT_CLOSE, self.on_close)
self.control_panel = Controls(self)
self.status_window = StatusWindow(self, textwidth=256)
hbox = wx.BoxSizer(wx.HORIZONTAL)
self.record_slider = wx.Slider(self,
value=0,
minValue=0,
maxValue=0,
style=wx.SL_HORIZONTAL | wx.SL_LABELS)
self.record_slider.Bind(wx.EVT_SLIDER, self.on_record_slider_scroll)
self.record_slider.Disable()
vbox = wx.BoxSizer(wx.VERTICAL)
vbox.Add(self.plotter, 1, wx.EXPAND)
vbox.Add(self.record_slider, 0, wx.ALL | wx.EXPAND, border=5)
self.left_vbox = vbox
hbox.Add(self.status_window, 0, wx.EXPAND)
hbox.Add(vbox, 1, wx.EXPAND)
hbox.Add(self.control_panel)
self.SetSizer(hbox)
self.control_panel.start_sim_button.Bind(wx.EVT_BUTTON,
self.cmd_simulate)
self.control_panel.stop_sim_button.Bind(wx.EVT_BUTTON,
self.cmd_stop_simulation)
self.control_panel.stop_sim_button.Disable()
self.control_panel.set_start_pos_button.Bind(
wx.EVT_BUTTON, lambda e: self.set_state(SimGui.StateSetStartPos))
self.control_panel.set_goal_pos_button.Bind(
wx.EVT_BUTTON, lambda e: self.set_state(SimGui.StateSetGoalPos))
self.control_panel.heading_text.SetValue('0.0')
self.control_panel.max_steps_text.SetValue('2000')
self.control_panel.step_size_text.SetValue('0.1')
self.control_panel.show_waypoints.SetValue(True)
self.control_panel.show_waypoints.Bind(
wx.EVT_CHECKBOX, lambda e: self.cmd_show_waypoints())
self.plotter.canvas.mpl_connect('button_press_event',
lambda ev: self.on_click_canvas(ev))
self.state = None
# Simulation related
self.sim_timer = wx.Timer(self)
self.Bind(wx.EVT_TIMER, self.on_timer, self.sim_timer)
self.sim_steps = 0
self.sim_state = SimGui.SimStateStopped
self.sim_max_steps = 0
self.last_start_pos = None
self.last_goal_pos = None
self.waypoints = []
self.last_waypoints = None
self.destination = None
self.reset_simulation()
self.sim_thread = None
self.agents = {}
# graph elements
self.cm = matplotlib.cm.get_cmap('Dark2')
self.h_start_point = None
self.h_goal_point = None
self.h_waypoints = None
self.h_legend = None
self.h_background_rect = None # A white patch to cover the whole area
self.background = None
self.background_bbox = None
self.background_ax_limit = None
def load_map(self, map):
self.map = map
self.map_visualizer = FindMapVisualizer(self.map)(self.map,
self.plotter.ax)
self.plotter.ax.set_title(map.__repr__(), fontsize=16)
self.plotter.canvas.draw()
self.map_visualizer.draw_map()
def reset_simulation(self):
self.waypoints = []
self.destination = None
self.sim_steps = 0
def set_state(self, state):
self.state = state
def draw_start_pos(self, x=None, y=None):
if x is None or y is None:
x, y = self.get_start_pos()
if self.h_start_point:
self.h_start_point.set_xdata(x)
self.h_start_point.set_ydata(y)
else:
self.h_start_point, = self.plotter.ax.plot(x,
y,
'o',
markeredgecolor=(0.8, 0,
0),
markerfacecolor='none')
self.plotter.canvas.draw()
def draw_goal_pos(self, x=None, y=None):
if x is None or y is None:
x, y = self.get_goal_pos()
if self.h_goal_point:
self.h_goal_point.set_xdata(x)
self.h_goal_point.set_ydata(y)
else:
self.h_goal_point, = self.plotter.ax.plot(x,
y,
'*',
color=(0.8, 0, 0))
self.plotter.canvas.draw()
def get_start_pos(self):
toks = self.control_panel.start_pos_text.GetValue().split(',')
return [float(t) for t in toks]
def get_goal_pos(self):
toks = self.control_panel.goal_pos_text.GetValue().split(',')
return [float(t) for t in toks]
def set_start_pos(self, x, y):
self.control_panel.start_pos_text.SetValue('%.2f, %.2f' % (x, y))
self.draw_start_pos(x, y)
def set_goal_pos(self, x, y):
self.control_panel.goal_pos_text.SetValue('%.2f, %.2f' % (x, y))
self.draw_goal_pos(x, y)
def set_heading(self, heading):
self.control_panel.heading_text.SetValue('%f' % heading)
def draw_waypoints(self, xs, ys):
if self.h_waypoints:
self.h_waypoints.set_xdata(xs)
self.h_waypoints.set_ydata(ys)
return
self.h_waypoints, = self.plotter.ax.plot(xs, ys, color='g', alpha=0.2)
def set_sim_state(self, state):
if state == SimGui.SimStateRunning:
self.control_panel.start_sim_button.SetLabel('Pause')
self.control_panel.stop_sim_button.Enable()
elif state == SimGui.SimStatePaused:
self.control_panel.start_sim_button.SetLabel('Resume')
self.control_panel.stop_sim_button.Enable()
elif state == SimGui.SimStateStopped:
self.control_panel.start_sim_button.SetLabel('Start')
self.control_panel.stop_sim_button.Disable()
self.sim_state = state
def cmd_stop_simulation(self, e):
self.stop()
print('simulation stopped')
def _finish_current_step(self):
if self.sim_thread is not None:
self.sim_thread.join()
self.sim_thread = None
start_time = time.time()
self.simulate_render(self.sim_steps)
print('render time: %.3f' % (time.time() - start_time))
self.sim_steps += 1
def on_timer(self, e):
if self.sim_thread is not None:
if self.sim_thread.is_alive():
return
self.sim_thread = None
start_time = time.time()
self.simulate_render(self.sim_steps)
print('render time: %.3f' % (time.time() - start_time))
if self.control_panel.save_screenshot.GetValue():
buf = self.make_screenshot()
cv2.imwrite('sim_images/%05d.png' % self.sim_steps, buf)
for name in self.agents:
agent = self.agents[name]['agent']
if hasattr(agent, 'img'):
img = self._prepare_img(agent.img, 256, 0.1)
cv2.imwrite(
'sim_images/%s_%05d.png' % (name, self.sim_steps),
img)
if self.sim_steps == 0:
# Call Layout() after rendering the first frame because the plot view may change
# its size.
self.Layout()
self.sim_steps += 1
n_sim_completed = 0
for name in self.agents:
agent = self.agents[name]['agent']
if agent.stopped():
print('agent %s stopped' % name)
n_sim_completed += 1
elif agent.reached_goal():
print('agent %s reached goal' % name)
n_sim_completed += 1
if n_sim_completed == len(self.agents):
print('simulation completes')
self.stop()
return
if self.sim_steps >= self.sim_max_steps:
print('max steps reached')
self.stop()
return
self.sim_thread = threading.Thread(target=lambda: self.simulate_step())
self.sim_thread.start()
def on_record_slider_scroll(self, e):
step = self.record_slider.GetValue()
for name in self.agents:
history = self.agents[name]['history']
state = history[min(step, len(history) - 1)]
self.agents[name]['agent'].restore_state(state)
self.simulate_render(step)
def pause(self):
self.sim_timer.Stop()
self._finish_current_step()
self.set_sim_state(SimGui.SimStatePaused)
self.record_slider.Enable()
self.record_slider.SetMin(0)
self.record_slider.SetMax(self.sim_steps - 1)
self.record_slider.SetValue(self.sim_steps - 1)
def resume(self):
self.set_sim_state(SimGui.SimStateRunning)
self.record_slider.Disable()
for name in self.agents:
history = self.agents[name]['history']
state = history[-1]
self.agents[name]['agent'].restore_state(state)
self.sim_timer.Start(1, oneShot=wx.TIMER_CONTINUOUS)
def stop(self):
if self.sim_state == SimGui.SimStateRunning:
self.pause()
self.reset_simulation()
self.sim_timer.Stop()
self.set_sim_state(SimGui.SimStateStopped)
self.record_slider.Disable()
def cmd_simulate(self, e):
if self.sim_state == SimGui.SimStateRunning:
self.pause()
return
elif self.sim_state == SimGui.SimStatePaused:
self.resume()
return
self.reset_simulation()
param_utils.global_params.reload()
start_pos = np.array(self.get_start_pos())
goal_pos = np.array(self.get_goal_pos())
if self.control_panel.no_goal.GetValue():
waypoints = []
elif self.control_panel.no_planning.GetValue():
waypoints = [goal_pos]
else:
if self.last_start_pos is not None and self.last_goal_pos is not None \
and np.all(self.last_start_pos == start_pos) and np.all(self.last_goal_pos == goal_pos):
print('reuse waypoints')
waypoints = self.last_waypoints
else:
if self.control_panel.goal_destination_mode.GetValue():
dest = tuple(
self.map.get_destination_from_map_coord(
goal_pos[0], goal_pos[1]))
waypoints = self.map.find_path_destination(start_pos, dest)
else:
waypoints = self.map.find_path(start_pos, goal_pos)
self.last_waypoints = waypoints
self.last_start_pos = start_pos
self.last_goal_pos = goal_pos
if waypoints is None:
print('infeasible')
return
self.init_agents()
if self.control_panel.replan.GetValue():
# TODO: remove code duplication
dest = tuple(
self.map.get_destination_from_map_coord(
goal_pos[0], goal_pos[1]))
if dest not in self.map.reachable_locs_per_destination:
print('invalid destination')
self.destination = dest
for name in self.agents:
agent = self.agents[name]['agent']
agent.reset()
agent.set_map(self.map)
agent.set_pos(start_pos)
agent.set_heading(
np.deg2rad(float(self.control_panel.heading_text.GetValue())))
agent.set_waypoints(waypoints)
if hasattr(agent, 'set_destination'):
# This agent uses destination as its goal
dest = tuple(
self.map.get_destination_from_map_coord(
goal_pos[0], goal_pos[1]))
if dest not in self.map.reachable_locs_per_destination:
print('invalid destination')
else:
agent.set_destination(dest)
print('destination label', dest)
if waypoints is not None and len(waypoints) > 0:
self.draw_waypoints(*zip(*waypoints))
self.waypoints = waypoints
self.draw_start_pos()
self.draw_goal_pos()
show_ticks = self.control_panel.show_axes_ticks.GetValue()
self.plotter.ax.get_xaxis().set_visible(show_ticks)
self.plotter.ax.get_yaxis().set_visible(show_ticks)
if self.control_panel.show_title.GetValue():
self.plotter.ax.set_title(self.map.__repr__(), fontsize=16)
else:
self.plotter.ax.set_title('')
agent_names = sorted(list(self.agents.keys()))
def to_uint8(color):
return int(color[0] * 255), int(color[1] * 255), int(color[2] *
255)
agent_colors = [to_uint8(self.agents[_]['color']) for _ in agent_names]
self.status_window.reset(agent_names, agent_colors)
self.sim_max_steps = int(self.control_panel.max_steps_text.GetValue())
self.sim_timer.Start(1, oneShot=wx.TIMER_CONTINUOUS)
self.set_sim_state(SimGui.SimStateRunning)
def _get_color(self, idx):
return self.cm(idx % self.cm.N)
def init_agents(self):
# Clear current agents' visualization items.
for _, agent in iteritems(self.agents):
agent['visualizer'].clear()
agent_switches = [
b.GetValue() for b in self.control_panel.agent_checkedbuttons
]
new_agents = {}
agent_idx = 0
for i in range(len(agent_switches)):
if not agent_switches[i]:
continue
agent_constructor = agent_factory.public_agents[i]
name = agent_constructor.__name__
# Reuse agent object if possible because some agents
# take a long time to initialize.
# FIXME: HACK!
# We don't reuse non-visual agents because we might change their parameters
# if name in self.agents and hasattr(self.agents[name]['agent'], 'img'):
# instance = self.agents[name]['agent']
# else:
agent_kwargs = str_to_dict(
self.control_panel.agent_preference_texts[i])
instance = agent_constructor(**agent_kwargs)
c = self._get_color(agent_idx)
vis = FindAgentVisualizer(instance)(self.plotter.ax,
draw_control_point=False,
traj_color=c,
obstacle_color=c,
heading_color=c,
active_wp_color=c,
goal_color=c,
label=name)
new_agents[name] = {
'agent': instance,
'color': c,
'visualizer': vis,
'legend_handler': AgentLegendHandler(c),
'history': [],
'sim_state': {
'traj': [],
'collisions': 0
}
}
agent_idx += 1
# Delete agents that are no longer used
for name, agent in iteritems(self.agents):
if name not in new_agents:
del agent['agent']
self.agents = new_agents
def draw_background(self, force_redraw=False):
'''
:return: draw background and cache it.
'''
# Pan/zoom can change axis limits. In that case we need to redraw the background.
x1, x2 = self.plotter.ax.get_xlim()
y1, y2 = self.plotter.ax.get_ylim()
def draw_legends():
agent_ids = self.agents.keys()
if self.control_panel.show_legends.GetValue():
self.h_legend = self.plotter.ax.legend(
agent_ids,
agent_ids,
handler_map={
agent_id: self.agents[agent_id]['legend_handler']
for agent_id in agent_ids
},
loc='lower left',
bbox_to_anchor=(0, -0.1),
ncol=5)
elif self.h_legend:
self.h_legend.remove()
self.h_legend = None
if self.background is None or \
self.background_bbox != self.plotter.ax.bbox.__repr__() or \
self.background_ax_limit != (x1, x2, y1, y2) or\
force_redraw:
self.plotter.ax.autoscale(False)
draw_legends()
# Cover all graphical elements
if self.h_background_rect is None:
self.h_background_rect = Rectangle((x1, y1),
x2 - x1,
y2 - y1,
color='w')
self.h_background_rect.set_zorder(10**5)
self.plotter.ax.add_patch(self.h_background_rect)
else:
self.h_background_rect.set_bounds(x1, y1, x2 - x1, y2 - y1)
self.h_background_rect.set_visible(True)
self.plotter.ax.draw_artist(self.h_background_rect)
self.plotter.canvas.draw()
# Draw the map
self.map_visualizer.draw_map()
self.h_background_rect.set_visible(False)
self.plotter.canvas.blit(self.plotter.ax.bbox)
self.background = self.plotter.canvas.copy_from_bbox(
self.plotter.ax.bbox)
self.background_bbox = self.plotter.ax.bbox.__repr__()
# limits might get changed. Retrieve new limits here.
x1, x2 = self.plotter.ax.get_xlim()
y1, y2 = self.plotter.ax.get_ylim()
self.background_ax_limit = (x1, x2, y1, y2)
else:
self.plotter.canvas.restore_region(self.background)
def simulate_step(self):
# Note that this runs in another thread.
step_size = float(self.control_panel.step_size_text.GetValue())
print('step %d step_size: %.3f' % (self.sim_steps, step_size))
# Update agent states
for name in self.agents:
start_time = time.time()
agent = self.agents[name]['agent']
if agent.collide(
) and self.control_panel.stop_on_collision.GetValue():
pass
else:
if self.control_panel.replan.GetValue():
waypoints = self.map.find_path_destination(
agent.pos, self.destination)
if waypoints is None or len(waypoints) == 0:
print('no valid waypoint. Use previous.')
else:
agent.set_waypoints(waypoints)
agent.wp_idx = 0
agent.step(step_size)
if agent.collide():
self.agents[name]['sim_state']['collisions'] += 1
self.agents[name]['history'].append(agent.save_state())
print('agent %s step time: %.3f '
'pos %.3f %.3f heading %.3f deg vel_norm %.3f collisions: %d' % \
(name, time.time() - start_time,
agent.pos[0], agent.pos[1],
np.rad2deg(agent.heading),
np.linalg.norm(agent.velocity),
self.agents[name]['sim_state']['collisions']))
pos = agent.pos
self.agents[name]['sim_state']['traj'].append([pos[0], pos[1]])
def simulate_render(self, step_idx):
# Always redraw at step 0
self.draw_background(force_redraw=(step_idx == 0))
self.plotter.ax.draw_artist(self.h_start_point)
self.plotter.ax.draw_artist(self.h_goal_point)
if self.h_waypoints is not None:
self.plotter.ax.draw_artist(self.h_waypoints)
for name in self.agents:
v = self.agents[name]
agent = v['agent']
sim_state = v['sim_state']
visualizer = v['visualizer']
visualizer.draw_trajectory(*zip(*sim_state['traj'][:step_idx + 1]))
visualizer.draw_agent_state(agent)
visualizer.set_show_trajectory(
self.control_panel.show_traj.GetValue())
visualizer.set_show_accel(self.control_panel.show_accel.GetValue())
visualizer.set_show_control_point_accels(
self.control_panel.show_control_point_accels.GetValue())
visualizer.set_show_control_point_metrics(
self.control_panel.show_control_point_metrics.GetValue())
visualizer.set_show_obstacles(
self.control_panel.show_obstacles.GetValue())
if hasattr(agent, 'img') and agent.img is not None:
img = (np.array(agent.img.transpose(1, 2, 0), copy=True) *
255.0).astype(np.uint8)
self._draw_all_waypoints(img, agent)
self.status_window.set_image(name, img)
self.status_window.set_text(
name, 'velocity', '%.3f m/s' % np.linalg.norm(agent.velocity))
self.status_window.set_text(
name, 'angular_vel',
'%.1f deg/s' % np.rad2deg(agent.angular_velocity))
self.plotter.canvas.blit(self.plotter.ax.bbox)
def _prepare_img(self, img, img_width, brightness_adj):
img = np.transpose(img, (1, 2, 0))
height = int(float(img_width) / img.shape[1] * img.shape[0])
img = cv2.resize(img, (img_width, height))
if img.dtype == np.float32:
img = np.clip(((img + brightness_adj) * 255), 0.0,
255.0).astype(np.uint8)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
return img
def _draw_all_waypoints(self, img, agent):
"""
Draw the groundtruth and neural waypoints if available
TODO: remove hardcoded colors.
:param img: 3-channel uint8 image
"""
if len(agent.goals_global) > 0:
self._draw_waypoint(img, agent,
agent.global_to_local(agent.goals_global[0]),
(0, 255, 0))
def _draw_waypoint(self, img, agent, wp_local, color):
if hasattr(agent, 'camera_pos'):
cam_pos = agent.camera_pos
else:
return
# Select a sensor that has 'h_fov' attribute.
for sensor in agent.sensors:
if hasattr(sensor, 'h_fov'):
fov = sensor.h_fov
img_width = img.shape[1]
def to_image_space(p):
size = max(int(10.0 / (np.linalg.norm(p) + 0.01)), 5)
l = img_width * 0.5 / np.tan(0.5 * fov)
d = p[1] / p[0] * l
return d, size
wp_local[0] -= cam_pos[0]
wp_local[1] -= cam_pos[1]
if wp_local[0] < 0:
# Waypoint behind the agent. Ignore.
return
d, s = to_image_space(wp_local)
cv2.circle(img, (img.shape[1] // 2 - int(d), img.shape[0] // 2), s,
color, 2)
def make_screenshot(self):
w, h = self.plotter.canvas.get_width_height()
buf = np.fromstring(self.plotter.canvas.tostring_rgb(), dtype=np.uint8)
buf = buf.reshape((h, w, 3))
buf = cv2.cvtColor(buf, cv2.COLOR_RGB2BGR)
# If the agent uses visual images, also save them.
img_width = 512
brightness_adj = 0.1
strip = np.ones((h, img_width, 3), np.uint8) * 128
y = 0
for name in self.agents:
agent = self.agents[name]['agent']
if hasattr(agent, 'img'):
img = self._prepare_img(agent.img, img_width, brightness_adj)
self._draw_all_waypoints(img, agent)
if hasattr(agent, 'img_left'):
img_left = self._prepare_img(agent.img_left, img_width,
brightness_adj)
img = np.concatenate([img_left, img], axis=1)
if hasattr(agent, 'img_right'):
img_right = self._prepare_img(agent.img_right, img_width,
brightness_adj)
img = np.concatenate([img, img_right], axis=1)
label_banner_cairo = np.zeros((32, img.shape[1]), np.uint32)
surface = cairo.ImageSurface.create_for_data(
label_banner_cairo, cairo.FORMAT_ARGB32,
label_banner_cairo.shape[1], label_banner_cairo.shape[0])
cr = cairo.Context(surface)
if agent.collide():
cr.set_source_rgb(255, 0, 0)
else:
cr.set_source_rgb(255, 255, 255)
cr.select_font_face('Helvetica', cairo.FONT_SLANT_NORMAL,
cairo.FONT_WEIGHT_NORMAL)
cr.set_font_size(24)
cr.move_to(0, 24)
cr.show_text(name)
label_banner = cairo_argb_to_opencv_bgr(label_banner_cairo)
img = np.concatenate([label_banner, img], axis=0)
if y + img.shape[0] > strip.shape[0]:
# Create new if current is full
buf = np.concatenate([buf, strip], axis=1)
strip = np.zeros((h, img_width, 3), np.uint8)
y = 0
# Enlarge the strip's width if img's width is larger than the strip's width
if strip.shape[1] < img.shape[1]:
padding = np.zeros(
(strip.shape[0], img.shape[1] - strip.shape[1],
strip.shape[2]),
dtype=strip.dtype)
strip = np.concatenate([strip, padding], axis=1)
strip[y:y + img.shape[0], 0:img.shape[1], :] = img
y += img.shape[0]
if y > 0:
buf = np.concatenate([buf, strip], axis=1)
return buf
def on_click_canvas(self, event):
if self.state == SimGui.StateSetStartPos:
self.set_start_pos(event.xdata, event.ydata)
self.state = SimGui.StateNone
elif self.state == SimGui.StateSetGoalPos:
self.set_goal_pos(event.xdata, event.ydata)
self.state = SimGui.StateNone
def on_close(self, e):
self.sim_timer.Stop()
del self.agents
e.Skip()
def cmd_show_waypoints(self):
if self.h_waypoints:
self.h_waypoints.set_visible(
self.control_panel.show_waypoints.GetValue())
class SimRenderer(object):
def __init__(self, map, ax, canvas):
self.ax = ax
self.canvas = canvas
self.map = map
self.map_visualizer = FindMapVisualizer(self.map)(self.map, self.ax)
self.agents = {}
self.show_legends = True
self.h_legend = None
self.background = None
self.background_ax_limit = None
self.background_bbox = None
self.h_background_rect = None
self.plotter = CachedPlotter(ax)
def set_agents(self, agents):
self.agents = agents
def clear(self):
self.plotter.clear()
def draw_background(self, force_redraw=False):
'''
:return: draw background and cache it.
'''
# Pan/zoom can change axis limits. In that case we need to redraw the background.
x1, x2 = self.ax.get_xlim()
y1, y2 = self.ax.get_ylim()
if self.background is None or \
self.background_bbox != self.ax.bbox.__repr__() or \
self.background_ax_limit != (x1, x2, y1, y2) or \
force_redraw:
self.ax.autoscale(False)
# Cover all graphical elements
if self.h_background_rect is None:
self.h_background_rect = Rectangle((x1, y1),
x2 - x1,
y2 - y1,
color='w')
self.h_background_rect.set_zorder(10**5)
self.ax.add_patch(self.h_background_rect)
else:
self.h_background_rect.set_bounds(x1, y1, x2 - x1, y2 - y1)
self.h_background_rect.set_visible(True)
self.ax.draw_artist(self.h_background_rect)
self.canvas.draw()
self.map_visualizer.draw_map()
self.h_background_rect.set_visible(False)
self.canvas.blit(self.ax.bbox)
self.background = self.canvas.copy_from_bbox(self.ax.bbox)
self.background_bbox = self.ax.bbox.__repr__()
# limits might get changed. Retrieve new limits here.
x1, x2 = self.ax.get_xlim()
y1, y2 = self.ax.get_ylim()
self.background_ax_limit = (x1, x2, y1, y2)
else:
self.canvas.restore_region(self.background)
def set_limits(self, x1, x2, y1, y2):
self.ax.set_xlim(x1, x2)
self.ax.set_ylim(y1, y2)
def get_limits(self):
x1, x2 = self.ax.get_xlim()
y1, y2 = self.ax.get_ylim()
return x1, x2, y1, y2
def set_viewport(self, x, y, scale):
"""
:param x, y: center of the viewport
:param scale: zoom factor w.r.t current viewport
"""
x_min, x_max, y_min, y_max = self.map.visible_map_bbox
w = (x_max - x_min) * scale
h = (y_max - y_min) * scale
xx1, xx2 = x - w / 2., x + w / 2.
yy1, yy2 = y - h / 2., y + h / 2.
self.ax.set_xlim(xx1, xx2)
self.ax.set_ylim(yy1, yy2)
def reset_viewport(self):
x_min, x_max, y_min, y_max = self.map.visible_map_bbox
self.ax.set_xlim([x_min, x_max])
self.ax.set_ylim([y_min, y_max])
def draw_agents(self):
for name in self.agents:
v = self.agents[name]
agent = v['agent']
visualizer = v['visualizer']
visualizer.draw_agent_state(agent)
def render(self, force_redraw=False, blit=True):
self.draw_background(force_redraw)
self.draw_agents()
if blit:
self.canvas.blit(self.ax.bbox)
def blit(self):
self.canvas.blit(self.ax.bbox)
def get_image(self):
w, h = self.canvas.get_width_height()
buf = np.fromstring(self.canvas.tostring_rgb(), dtype=np.uint8)
buf = buf.reshape((h, w, 3))
return buf
def save(self, filename, **kwargs):
self.canvas.print_figure(filename, **kwargs)
'-s',
f'{width}x{height}', # size of image string
'-f',
'rawvideo',
'-pix_fmt',
'rgba', # format
'-i',
'-', # tell ffmpeg to expect raw video from the pipe
'-c:v',
'hevc_videotoolbox',
'-tag:v',
'hvc1',
'-profile:v',
'main10',
'-b:v',
'16M',
OUTPUT_VIDEO,
)
ffmpeg = subprocess.Popen(cmd, stdin=subprocess.PIPE, bufsize=10**8)
for i in range(40):
print('task.i:', i)
rect.set_bounds((0, 0, i, height))
fig.savefig(ffmpeg.stdin, format='rgba', dpi=100)
ffmpeg.communicate()
# out, err = ffmpeg.communicate()
subprocess.run(('open', OUTPUT_VIDEO))
class ThreeCompVisualisation:
"""
Basis to visualise power flow within the hydraulics model as an animation or simulation
"""
def __init__(self, agent: ThreeCompHydAgent,
axis: plt.axis = None,
animated: bool = False,
detail_annotations: bool = False,
basic_annotations: bool = True,
black_and_white: bool = False):
"""
Whole visualisation setup using given agent's parameters
:param agent: The agent to be visualised
:param axis: If set, the visualisation will be drawn using the provided axis object.
Useful for animations or to display multiple models in one plot
:param animated: If true the visualisation is set up to deal with frame updates.
See animation script for more details.
:param detail_annotations: If true, tank distances and sizes are annotated as well.
:param basic_annotations: If true, U, LF, and LS are visible
:param black_and_white: If true, the visualisation is in black and white
"""
# matplotlib fontsize
rcParams['font.size'] = 10
# plot if no axis was assigned
if axis is None:
fig = plt.figure(figsize=(8, 5))
self._ax1 = fig.add_subplot(1, 1, 1)
else:
fig = None
self._ax1 = axis
if black_and_white:
self.__u_color = (0.7, 0.7, 0.7)
self.__lf_color = (0.5, 0.5, 0.5)
self.__ls_color = (0.3, 0.3, 0.3)
self.__ann_color = (0, 0, 0)
self.__p_color = (0.5, 0.5, 0.5)
elif not black_and_white:
self.__u_color = "tab:cyan"
self.__lf_color = "tab:orange"
self.__ls_color = "tab:red"
self.__ann_color = "tab:blue"
self.__p_color = "tab:green"
# basic parameters for setup
self._animated = animated
self.__detail_annotations = detail_annotations
self._agent = agent
self.__offset = 0.2
# U tank with three stripes
self.__width_u = 0.3
self._u = None
self._u1 = None
self._u2 = None
self._r1 = None # line marking flow from U to LF
self._ann_u = None # U annotation
# LF tank
self._lf = None
self._h = None # fill state
self._ann_lf = None # annotation
# LS tank
self._ls = None
self._g = None
self._ann_ls = None # annotation LS
self._r2 = None # line marking flow from LS to LF
# finish the basic layout
self.__set_basic_layout()
self.update_basic_layout(agent)
# now the animation components
if self._animated:
# U flow
self._arr_u_flow = None
self._ann_u_flow = None
# flow out of tap
self._arr_power_flow = None
self._ann_power_flow = None
# LS flow (R2)
self._arr_r2_l_pos = None
self._arr_r2_flow = None
self._ann_r2_flow = None
# time information annotation
self._ann_time = None
self.__set_animation_layout()
self._ax1.add_artist(self._ann_time)
# basic annotations are U, LF, and LS
if not basic_annotations:
self.hide_basic_annotations()
# add layout for detailed annotations
# detail annotation add greek letters for distances and positions
if self.__detail_annotations:
self.__set_detailed_annotations_layout()
self._ax1.set_xlim(0, 1.05)
self._ax1.set_ylim(0, 1.2)
else:
self._ax1.set_xlim(0, 1.0)
self._ax1.set_ylim(0, 1.2)
if self.__detail_annotations and self._animated:
raise UserWarning("Detailed annotations and animation cannot be combined")
self._ax1.set_axis_off()
# display plot if no axis object was assigned
if fig is not None:
plt.subplots_adjust(left=0.02, bottom=0.02, right=0.98, top=0.98)
plt.show()
plt.close(fig)
def __set_detailed_annotations_layout(self):
"""
Adds components required for a detailed
annotations view with denoted positions and distances
"""
u_width = self.__width_u
ls_left = self._ls.get_x()
ls_width = self._ls.get_width()
lf_left = self._lf.get_x()
lf_width = self._lf.get_width()
ls_height = self._ls.get_height()
# some offset to the bottom
offset = self.__offset
phi_o = self._agent.phi + offset
gamma_o = self._agent.gamma + offset
rcParams['text.usetex'] = True
self._ann_u_flow = Text(text="$M_{U}$", ha='right', fontsize="xx-large", x=u_width + 0.09,
y=phi_o - 0.08)
ann_p_ae = Text(text="$p_{U}$", ha='right', fontsize="xx-large", x=u_width + 0.07,
y=phi_o + 0.03)
self._arr_u_flow = FancyArrowPatch((u_width, phi_o),
(u_width + 0.1, phi_o),
arrowstyle='-|>',
mutation_scale=30,
lw=2,
color=self.__u_color)
self._ax1.annotate('$\phi$',
xy=(u_width / 2, phi_o),
xytext=(u_width / 2, (phi_o - offset) / 2 + offset - 0.015),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.annotate('$\phi$',
xy=(u_width / 2, offset),
xytext=(u_width / 2, (phi_o - offset) / 2 + offset - 0.015),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ann_power_flow = Text(text="$p$", ha='center', fontsize="xx-large", x=self._ann_lf.get_position()[0],
y=offset - 0.06)
self._arr_power_flow = FancyArrowPatch((self._ann_lf.get_position()[0], offset - 0.078),
(self._ann_lf.get_position()[0], 0.0),
arrowstyle='-|>',
mutation_scale=30,
lw=2,
color=self.__p_color)
self._ax1.annotate('$h$',
xy=(self._ann_lf.get_position()[0] + 0.07, 1 + offset),
xytext=(self._ann_lf.get_position()[0] + 0.07, 1 + offset - 0.30),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.annotate('$h$',
xy=(self._ann_lf.get_position()[0] + 0.07, 1 + offset - 0.55),
xytext=(self._ann_lf.get_position()[0] + 0.07, 1 + offset - 0.30),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._h.update(dict(xy=(lf_left, offset),
width=lf_width,
height=1 - 0.55,
color=self.__lf_color))
self._ax1.annotate('$g$',
xy=(ls_left + ls_width + 0.02, ls_height + gamma_o),
xytext=(ls_left + ls_width + 0.02, ls_height * 0.61 + gamma_o),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.annotate('$g$',
xy=(ls_left + ls_width + 0.02, ls_height * 0.3 + gamma_o),
xytext=(ls_left + ls_width + 0.02, ls_height * 0.61 + gamma_o),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._g.update(dict(xy=(ls_left, gamma_o),
width=ls_width,
height=ls_height * 0.3,
color=self.__ls_color))
ann_p_an = Text(text="$p_{L}$", ha='left', usetex=True, fontsize="xx-large", x=ls_left - 0.06,
y=gamma_o + 0.11)
ann_arr_flow = Text(text="$M_{LS}$", ha='left', usetex=True, fontsize="xx-large", x=ls_left - 0.09,
y=gamma_o + 0.03)
self._ann_r2_flow = Text(text="$M_{LF}$", ha='left', usetex=True, fontsize="xx-large", x=ls_left - 0.04,
y=gamma_o - 0.07)
self._arr_r2_flow = FancyArrowPatch((ls_left - 0.1, gamma_o),
(ls_left, gamma_o),
arrowstyle='<|-|>',
mutation_scale=30,
lw=2,
color=self.__ls_color)
self._ax1.annotate('$\\theta$',
xy=(ls_left + ls_width / 2, 1 + offset),
xytext=(
ls_left + ls_width / 2,
1 - (1 - (ls_height + gamma_o - offset)) / 2 + offset - 0.015),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.annotate('$\\theta$',
xy=(ls_left + ls_width / 2, ls_height + gamma_o),
xytext=(
ls_left + ls_width / 2,
1 - (1 - (ls_height + gamma_o - offset)) / 2 + offset - 0.015),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.annotate('$\\gamma$',
xy=(ls_left + ls_width / 2, offset),
xytext=(ls_left + ls_width / 2, (gamma_o - offset) / 2 + offset - 0.015),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.annotate('$\\gamma$',
xy=(ls_left + ls_width / 2, gamma_o),
xytext=(ls_left + ls_width / 2, (gamma_o - offset) / 2 + offset - 0.015),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.annotate('$1$',
xy=(1.05, 0 + offset),
xytext=(1.05, 0.5 + offset),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.annotate('$1$',
xy=(1.05, 1 + offset),
xytext=(1.05, 0.5 + offset),
ha='center',
fontsize="xx-large",
arrowprops=dict(arrowstyle='-|>',
ls='-',
fc=self.__ann_color)
)
self._ax1.add_artist(ann_arr_flow)
# self._ax1.add_artist(ann_p_an)
# self._ax1.add_artist(ann_p_ae)
self._ax1.axhline(offset, linestyle='--', color=self.__ann_color)
self._ax1.axhline(1 + offset - 0.001, linestyle='--', color=self.__ann_color)
self._ax1.add_artist(self._ann_power_flow)
self._ax1.add_artist(self._arr_power_flow)
self._ax1.add_artist(self._arr_u_flow)
self._ax1.add_artist(self._ann_u_flow)
self._ax1.add_artist(self._arr_r2_flow)
self._ax1.add_artist(self._ann_r2_flow)
def __set_animation_layout(self):
"""
Adds layout components that are required for an animation
"""
offset = self.__offset
o_width = self.__width_u
phi_o = self._agent.phi + offset
gamma_o = self._agent.gamma + offset
# U flow (R1)
self._arr_u_flow = FancyArrowPatch((o_width, phi_o),
(o_width + 0.1, phi_o),
arrowstyle='simple',
mutation_scale=0,
ec='white',
fc=self.__u_color)
self._ann_u_flow = Text(text="flow: ", ha='right', fontsize="large", x=o_width, y=phi_o - 0.05)
# Tap flow (Power)
self._arr_power_flow = FancyArrowPatch((self._ann_lf.get_position()[0], offset - 0.05),
(self._ann_lf.get_position()[0], 0.0),
arrowstyle='simple',
mutation_scale=0,
ec='white',
color=self.__p_color)
self._ann_power_flow = Text(text="flow: ", ha='center', fontsize="large", x=self._ann_lf.get_position()[0],
y=offset - 0.05)
# LS flow (R2)
self._arr_r2_l_pos = [(self._ls.get_x(), gamma_o),
(self._ls.get_x() - 0.1, gamma_o)]
self._arr_r2_flow = FancyArrowPatch(self._arr_r2_l_pos[0],
self._arr_r2_l_pos[1],
arrowstyle='simple',
mutation_scale=0,
ec='white',
color=self.__ls_color)
self._ann_r2_flow = Text(text="flow: ", ha='left', fontsize="large", x=self._ls.get_x(),
y=gamma_o - 0.05)
# information annotation
self._ann_time = Text(x=1, y=0.9 + offset, ha="right")
self._ax1.add_artist(self._ann_power_flow)
self._ax1.add_artist(self._arr_power_flow)
self._ax1.add_artist(self._arr_u_flow)
self._ax1.add_artist(self._ann_u_flow)
self._ax1.add_artist(self._arr_r2_flow)
self._ax1.add_artist(self._ann_r2_flow)
def __set_basic_layout(self):
"""
updates position estimations and layout
"""
# get sizes from agent
lf = self._agent.lf
ls = self._agent.ls
ls_height = self._agent.height_ls
# u_left is 0
u_width = self.__width_u
# determine width with size ratio retained
lf_left = u_width + 0.1
lf_width = ((lf * ls_height) * (1 - lf_left - 0.1)) / ls
ls_left = lf_left + lf_width + 0.1
ls_width = 1 - ls_left
# some offset to the bottom
offset = self.__offset
phi_o = self._agent.phi + offset
gamma_o = self._agent.gamma + offset
# S tank
self._u = Rectangle((0.0, phi_o), 0.05, 1 - self._agent.phi, color=self.__u_color, alpha=0.3)
self._u1 = Rectangle((0.05, phi_o), 0.05, 1 - self._agent.phi, color=self.__u_color, alpha=0.6)
self._u2 = Rectangle((0.1, phi_o), u_width - 0.1, 1 - self._agent.phi, color=self.__u_color)
self._r1 = Line2D([u_width, u_width + 0.1],
[phi_o, phi_o],
color=self.__u_color)
self._ann_u = Text(text="$U$", ha='center', fontsize="xx-large",
x=u_width / 2,
y=((1 - self._agent.phi) / 2) + phi_o - 0.02)
# LF vessel
self._lf = Rectangle((lf_left, offset), lf_width, 1, fill=False, ec="black")
self._h = Rectangle((lf_left, offset), lf_width, 1, color=self.__lf_color)
self._ann_lf = Text(text="$LF$", ha='center', fontsize="xx-large",
x=lf_left + (lf_width / 2),
y=offset + 0.5 - 0.02)
# LS vessel
self._ls = Rectangle((ls_left, gamma_o), ls_width, ls_height, fill=False, ec="black")
self._g = Rectangle((ls_left, gamma_o), ls_width, ls_height, color=self.__ls_color)
self._r2 = Line2D([ls_left, ls_left - 0.1],
[gamma_o, gamma_o],
color=self.__ls_color)
self._ann_ls = Text(text="$LS$", ha='center', fontsize="xx-large",
x=ls_left + (ls_width / 2),
y=gamma_o + (ls_height / 2) - 0.02)
# the basic layout
self._ax1.add_line(self._r1)
self._ax1.add_line(self._r2)
self._ax1.add_artist(self._u)
self._ax1.add_artist(self._u1)
self._ax1.add_artist(self._u2)
self._ax1.add_artist(self._lf)
self._ax1.add_artist(self._ls)
self._ax1.add_artist(self._h)
self._ax1.add_artist(self._g)
self._ax1.add_artist(self._ann_u)
self._ax1.add_artist(self._ann_lf)
self._ax1.add_artist(self._ann_ls)
def update_basic_layout(self, agent):
"""
updates tank positions and sizes according to new agent
:param agent: agent to be visualised
"""
self._agent = agent
# get sizes from agent
lf = agent.lf
ls = agent.ls
ls_heigh = agent.height_ls
# o_left is 0
u_width = self.__width_u
# determine width with size ratio retained
lf_left = u_width + 0.1
lf_width = ((lf * ls_heigh) * (1 - lf_left - 0.1)) / (ls + lf * ls_heigh)
ls_left = lf_left + lf_width + 0.1
ls_width = 1 - ls_left
# some offset to the bottom
offset = self.__offset
phi_o = agent.phi + offset
gamma_o = agent.gamma + offset
# S tank
self._u.set_bounds(0.0, phi_o, 0.05, 1 - self._agent.phi)
self._u1.set_bounds(0.05, phi_o, 0.05, 1 - self._agent.phi)
self._u2.set_bounds(0.1, phi_o, u_width - 0.1, 1 - self._agent.phi)
self._r1.set_xdata([u_width, u_width + 0.1])
self._r1.set_ydata([phi_o, phi_o])
self._ann_u.set_position(xy=(u_width / 2, ((1 - self._agent.phi) / 2) + phi_o - 0.02))
# LF vessel
self._lf.set_bounds(lf_left, offset, lf_width, 1)
self._h.set_bounds(lf_left, offset, lf_width, 1)
self._ann_lf.set_position(xy=(lf_left + (lf_width / 2),
offset + 0.5 - 0.02))
# LS vessel
self._ls.set_bounds(ls_left, gamma_o, ls_width, ls_heigh)
self._g.set_bounds(ls_left, gamma_o, ls_width, ls_heigh)
self._r2.set_xdata([ls_left, ls_left - 0.1])
self._r2.set_ydata([gamma_o, gamma_o])
self._ann_ls.set_position(xy=(ls_left + (ls_width / 2), gamma_o + (ls_heigh / 2) - 0.02))
# update levels
self._h.set_height(1 - self._agent.get_h())
self._g.set_height(self._agent.height_ls - self._agent.get_g())
def hide_basic_annotations(self):
"""
Simply hides the S, LF, and LS text
"""
self._ann_u.set_text("")
self._ann_lf.set_text("")
self._ann_ls.set_text("")
def update_animation_data(self, frame_number):
"""
For animations and simulations.
The function to call at each frame.
:param frame_number: frame number has to be taken because of parent class method
:return: an iterable of artists
"""
if not self._animated:
raise UserWarning("Animation flag has to be enabled in order to use this function")
# perform one step
cur_time = self._agent.get_time()
power = self._agent.perform_one_step()
# draw some information
self._ann_time.set_text("agent \n time: {}".format(int(cur_time)))
# power arrow
self._ann_power_flow.set_text("power: {}".format(round(power)))
self._arr_power_flow.set_mutation_scale(math.log(power + 1) * 10)
# oxygen arrow
p_u = round(self._agent.get_p_u() * self._agent.hz, 1)
max_str = "(MAX)" if p_u == self._agent.m_u else ""
self._ann_u_flow.set_text("flow: {} {}".format(p_u, max_str))
self._arr_u_flow.set_mutation_scale(math.log(p_u + 1) * 10)
# lactate arrow
p_g = round(self._agent.get_p_l() * self._agent.hz, 1)
if p_g < 0:
max_str = "(MAX)" if p_g == self._agent.m_lf else ""
self._arr_r2_flow.set_positions(self._arr_r2_l_pos[1], self._arr_r2_l_pos[0])
else:
max_str = "(MAX)" if p_g == self._agent.m_ls else ""
self._arr_r2_flow.set_positions(self._arr_r2_l_pos[0], self._arr_r2_l_pos[1])
self._ann_r2_flow.set_text("flow: {} {}".format(p_g, max_str))
self._arr_r2_flow.set_mutation_scale(math.log(abs(p_g) + 1) * 10)
# update levels
self._h.set_height(1 - self._agent.get_h())
self._g.set_height(self._agent.height_ls - self._agent.get_g())
# list of artists to be drawn
return [self._ann_time,
self._ann_power_flow,
self._arr_power_flow,
self._g,
self._h]
