# -*- coding: utf-8 -*- """Shared GUI classes and functions.""" # Authors: Alex Rockhill # # License: BSD (3-clause) import os import os.path as op import numpy as np from functools import partial from qtpy import QtCore from qtpy.QtCore import Slot from qtpy.QtWidgets import (QMainWindow, QGridLayout, QVBoxLayout, QHBoxLayout, QLabel, QMessageBox, QWidget, QPlainTextEdit) from matplotlib import patheffects from matplotlib.backends.backend_qt5agg import FigureCanvas from matplotlib.figure import Figure from matplotlib.patches import Rectangle from .._freesurfer import _import_nibabel from ..viz.backends.renderer import _get_renderer from ..viz.utils import safe_event from ..surface import _read_mri_surface, _marching_cubes from ..transforms import apply_trans, _frame_to_str from ..utils import logger, _check_fname, verbose, warn, get_subjects_dir _IMG_LABELS = [['I', 'P'], ['I', 'L'], ['P', 'L']] _ZOOM_STEP_SIZE = 5 @verbose def _load_image(img, verbose=None): """Load data from a 3D image file (e.g. CT, MR).""" nib = _import_nibabel('use GUI') if not isinstance(img, nib.spatialimages.SpatialImage): logger.info(f'Loading {img}') _check_fname(img, overwrite='read', must_exist=True) img = nib.load(img) # get data orig_data = np.array(img.dataobj).astype(np.float32) # reorient data to RAS ornt = nib.orientations.axcodes2ornt( nib.orientations.aff2axcodes(img.affine)).astype(int) ras_ornt = nib.orientations.axcodes2ornt('RAS') ornt_trans = nib.orientations.ornt_transform(ornt, ras_ornt) img_data = nib.orientations.apply_orientation(orig_data, ornt_trans) orig_mgh = nib.MGHImage(orig_data, img.affine) aff_trans = nib.orientations.inv_ornt_aff(ornt_trans, img.shape) vox_ras_t = np.dot(orig_mgh.header.get_vox2ras_tkr(), aff_trans) return img_data, vox_ras_t def _make_slice_plot(width=4, height=4, dpi=300): fig = Figure(figsize=(width, height), dpi=dpi) canvas = FigureCanvas(fig) ax = fig.subplots() fig.subplots_adjust(bottom=0, left=0, right=1, top=1, wspace=0, hspace=0) ax.set_facecolor('k') # clean up excess plot text, invert ax.invert_yaxis() ax.set_xticks([]) ax.set_yticks([]) return canvas, fig class SliceBrowser(QMainWindow): """Navigate between slices of an MRI, CT, etc. image.""" _xy_idx = ( (1, 2), (0, 2), (0, 1), ) def __init__(self, base_image=None, subject=None, subjects_dir=None, verbose=None): """GUI for browsing slices of anatomical images.""" # initialize QMainWindow class super(SliceBrowser, self).__init__() self._verbose = verbose # if bad/None subject, will raise an informative error when loading MRI subject = os.environ.get('SUBJECT') if subject is None else subject subjects_dir = get_subjects_dir(subjects_dir, raise_error=True) self._subject_dir = op.join(subjects_dir, subject) self._load_image_data(base_image=base_image) # GUI design # Main plots: make one plot for each view; sagittal, coronal, axial self._plt_grid = QGridLayout() self._figs = list() for i in range(3): canvas, fig = _make_slice_plot() self._plt_grid.addWidget(canvas, i // 2, i % 2) self._figs.append(fig) self._renderer = _get_renderer( name='Slice Browser', size=(400, 400), bgcolor='w') self._plt_grid.addWidget(self._renderer.plotter, 1, 1) self._set_ras([0., 0., 0.], update_plots=False) self._plot_images() self._configure_ui() def _configure_ui(self): bottom_hbox = self._configure_status_bar() # Put everything together plot_ch_hbox = QHBoxLayout() plot_ch_hbox.addLayout(self._plt_grid) main_vbox = QVBoxLayout() main_vbox.addLayout(plot_ch_hbox) main_vbox.addLayout(bottom_hbox) central_widget = QWidget() central_widget.setLayout(main_vbox) self.setCentralWidget(central_widget) def _load_image_data(self, base_image=None): """Get image data to display and transforms to/from vox/RAS.""" # allows recon-all not to be finished (T1 made in a few minutes) mri_img = 'brain' if op.isfile(op.join( self._subject_dir, 'mri', 'brain.mgz')) else 'T1' self._mri_data, self._vox_ras_t = _load_image( op.join(self._subject_dir, 'mri', f'{mri_img}.mgz')) self._ras_vox_t = np.linalg.inv(self._vox_ras_t) self._voxel_sizes = np.array(self._mri_data.shape) # We need our extents to land the centers of each pixel on the voxel # number. This code assumes 1mm isotropic... img_delta = 0.5 self._img_extents = list( [-img_delta, self._voxel_sizes[idx[0]] - img_delta, -img_delta, self._voxel_sizes[idx[1]] - img_delta] for idx in self._xy_idx) # ready alternate base image if provided, otherwise use brain/T1 if base_image is None: self._base_data = self._mri_data else: self._base_data, vox_ras_t = _load_image(base_image) if self._mri_data.shape != self._base_data.shape or \ not np.allclose(self._vox_ras_t, vox_ras_t, rtol=1e-6): raise ValueError('Base image is not aligned to MRI, got ' f'Base shape={self._base_data.shape}, ' f'MRI shape={self._mri_data.shape}, ' f'Base affine={vox_ras_t} and ' f'MRI affine={self._vox_ras_t}') if op.exists(op.join(self._subject_dir, 'surf', 'lh.seghead')): self._head = _read_mri_surface( op.join(self._subject_dir, 'surf', 'lh.seghead')) assert _frame_to_str[self._head['coord_frame']] == 'mri' else: warn('`seghead` not found, using marching cubes on CT for ' 'head plot, use :ref:`mne.bem.make_scalp_surfaces` ' 'to add the scalp surface instead of skull from the CT') self._head = None if op.exists(op.join(self._subject_dir, 'surf', 'lh.pial')): self._lh = _read_mri_surface( op.join(self._subject_dir, 'surf', 'lh.pial')) assert _frame_to_str[self._lh['coord_frame']] == 'mri' self._rh = _read_mri_surface( op.join(self._subject_dir, 'surf', 'rh.pial')) assert _frame_to_str[self._rh['coord_frame']] == 'mri' else: warn('`pial` surface not found, skipping adding to 3D ' 'plot. This indicates the Freesurfer recon-all ' 'has not finished or has been modified and ' 'these files have been deleted.') self._lh = self._rh = None def _plot_images(self): """Use the MRI or CT to make plots.""" # Plot sagittal (0), coronal (1) or axial (2) view self._images = dict(base=list(), cursor_v=list(), cursor_h=list(), bounds=list()) img_min = np.nanmin(self._base_data) img_max = np.nanmax(self._base_data) text_kwargs = dict(fontsize='medium', weight='bold', color='#66CCEE', family='monospace', ha='center', va='center', path_effects=[patheffects.withStroke( linewidth=4, foreground="k", alpha=0.75)]) xyz = apply_trans(self._ras_vox_t, self._ras) for axis in range(3): plot_x_idx, plot_y_idx = self._xy_idx[axis] fig = self._figs[axis] ax = fig.axes[0] img_data = np.take(self._base_data, self._current_slice[axis], axis=axis).T self._images['base'].append(ax.imshow( img_data, cmap='gray', aspect='auto', zorder=1, vmin=img_min, vmax=img_max)) img_extent = self._img_extents[axis] # x0, x1, y0, y1 w, h = np.diff(np.array(img_extent).reshape(2, 2), axis=1)[:, 0] self._images['bounds'].append(Rectangle( img_extent[::2], w, h, edgecolor='w', facecolor='none', alpha=0.25, lw=0.5, zorder=1.5)) ax.add_patch(self._images['bounds'][-1]) v_x = (xyz[plot_x_idx],) * 2 v_y = img_extent[2:4] self._images['cursor_v'].append(ax.plot( v_x, v_y, color='lime', linewidth=0.5, alpha=0.5, zorder=8)[0]) h_y = (xyz[plot_y_idx],) * 2 h_x = img_extent[0:2] self._images['cursor_h'].append(ax.plot( h_x, h_y, color='lime', linewidth=0.5, alpha=0.5, zorder=8)[0]) # label axes self._figs[axis].text(0.5, 0.05, _IMG_LABELS[axis][0], **text_kwargs) self._figs[axis].text(0.05, 0.5, _IMG_LABELS[axis][1], **text_kwargs) self._figs[axis].axes[0].axis(img_extent) self._figs[axis].canvas.mpl_connect( 'scroll_event', self._on_scroll) self._figs[axis].canvas.mpl_connect( 'button_release_event', partial(self._on_click, axis=axis)) # add head and brain in mm (convert from m) if self._head is None: logger.info('Using marching cubes on CT for the ' '3D visualization panel') rr, tris = _marching_cubes(np.where( self._base_data < np.quantile(self._base_data, 0.95), 0, 1), [1])[0] rr = apply_trans(self._vox_ras_t, rr) self._renderer.mesh( *rr.T, triangles=tris, color='gray', opacity=0.2, reset_camera=False, render=False) else: self._renderer.mesh( *self._head['rr'].T * 1000, triangles=self._head['tris'], color='gray', opacity=0.2, reset_camera=False, render=False) if self._lh is not None and self._rh is not None: self._renderer.mesh( *self._lh['rr'].T * 1000, triangles=self._lh['tris'], color='white', opacity=0.2, reset_camera=False, render=False) self._renderer.mesh( *self._rh['rr'].T * 1000, triangles=self._rh['tris'], color='white', opacity=0.2, reset_camera=False, render=False) self._renderer.set_camera(azimuth=90, elevation=90, distance=300, focalpoint=tuple(self._ras)) # update plots self._draw() self._renderer._update() def _configure_status_bar(self, hbox=None): """Make a bar at the bottom with information in it.""" hbox = QHBoxLayout() if hbox is None else hbox self._intensity_label = QLabel('') # update later hbox.addWidget(self._intensity_label) VOX_label = QLabel('VOX =') self._VOX_textbox = QPlainTextEdit('') # update later self._VOX_textbox.setMaximumHeight(25) self._VOX_textbox.setMaximumWidth(125) self._VOX_textbox.focusOutEvent = self._update_VOX self._VOX_textbox.textChanged.connect(self._check_update_VOX) hbox.addWidget(VOX_label) hbox.addWidget(self._VOX_textbox) RAS_label = QLabel('RAS =') self._RAS_textbox = QPlainTextEdit('') # update later self._RAS_textbox.setMaximumHeight(25) self._RAS_textbox.setMaximumWidth(200) self._RAS_textbox.focusOutEvent = self._update_RAS self._RAS_textbox.textChanged.connect(self._check_update_RAS) hbox.addWidget(RAS_label) hbox.addWidget(self._RAS_textbox) self._update_moved() # update text now return hbox def _update_camera(self, render=False): """Update the camera position.""" self._renderer.set_camera( # needs fix, distance moves when focal point updates distance=self._renderer.plotter.camera.distance * 0.9, focalpoint=tuple(self._ras), reset_camera=False) def _on_scroll(self, event): """Process mouse scroll wheel event to zoom.""" self._zoom(event.step, draw=True) def _zoom(self, sign=1, draw=False): """Zoom in on the image.""" delta = _ZOOM_STEP_SIZE * sign for axis, fig in enumerate(self._figs): xmid = self._images['cursor_v'][axis].get_xdata()[0] ymid = self._images['cursor_h'][axis].get_ydata()[0] xmin, xmax = fig.axes[0].get_xlim() ymin, ymax = fig.axes[0].get_ylim() xwidth = (xmax - xmin) / 2 - delta ywidth = (ymax - ymin) / 2 - delta if xwidth <= 0 or ywidth <= 0: return fig.axes[0].set_xlim(xmid - xwidth, xmid + xwidth) fig.axes[0].set_ylim(ymid - ywidth, ymid + ywidth) if draw: self._figs[axis].canvas.draw() @Slot() def _update_RAS(self, event): """Interpret user input to the RAS textbox.""" text = self._RAS_textbox.toPlainText() ras = self._convert_text(text, 'ras') if ras is not None: self._set_ras(ras) @Slot() def _update_VOX(self, event): """Interpret user input to the RAS textbox.""" text = self._VOX_textbox.toPlainText() ras = self._convert_text(text, 'vox') if ras is not None: self._set_ras(ras) def _convert_text(self, text, text_kind): text = text.replace('\n', '') vals = text.split(',') if len(vals) != 3: vals = text.split(' ') # spaces also okay as in freesurfer vals = [var.lstrip().rstrip() for var in vals] try: vals = np.array([float(var) for var in vals]).reshape(3) except Exception: self._update_moved() # resets RAS label return if text_kind == 'vox': vox = vals ras = apply_trans(self._vox_ras_t, vox) else: assert text_kind == 'ras' ras = vals vox = apply_trans(self._ras_vox_t, ras) wrong_size = any(var < 0 or var > n - 1 for var, n in zip(vox, self._voxel_sizes)) if wrong_size: self._update_moved() # resets RAS label return return ras @property def _ras(self): return self._ras_safe def _set_ras(self, ras, update_plots=True): ras = np.asarray(ras, dtype=float) assert ras.shape == (3,) msg = ', '.join(f'{x:0.2f}' for x in ras) logger.debug(f'Trying RAS: ({msg}) mm') # clip to valid vox = apply_trans(self._ras_vox_t, ras) vox = np.array([ np.clip(d, 0, self._voxel_sizes[ii] - 1) for ii, d in enumerate(vox)]) # transform back, make write-only self._ras_safe = apply_trans(self._vox_ras_t, vox) self._ras_safe.flags['WRITEABLE'] = False msg = ', '.join(f'{x:0.2f}' for x in self._ras_safe) logger.debug(f'Setting RAS: ({msg}) mm') if update_plots: self._move_cursors_to_pos() @property def _vox(self): return apply_trans(self._ras_vox_t, self._ras) @property def _current_slice(self): return self._vox.round().astype(int) @Slot() def _check_update_RAS(self): """Check whether the RAS textbox is done being edited.""" if '\n' in self._RAS_textbox.toPlainText(): self._update_RAS(event=None) @Slot() def _check_update_VOX(self): """Check whether the VOX textbox is done being edited.""" if '\n' in self._VOX_textbox.toPlainText(): self._update_VOX(event=None) def _draw(self, axis=None): """Update the figures with a draw call.""" for axis in (range(3) if axis is None else [axis]): self._figs[axis].canvas.draw() def _update_base_images(self, axis=None, draw=False): """Update the base images.""" for axis in range(3) if axis is None else [axis]: img_data = np.take(self._base_data, self._current_slice[axis], axis=axis).T self._images['base'][axis].set_data(img_data) if draw: self._draw(axis) def _update_images(self, axis=None, draw=True): """Update CT and channel images when general changes happen.""" self._update_base_images(axis=axis) if draw: self._draw(axis) def _move_cursors_to_pos(self): """Move the cursors to a position.""" for axis in range(3): x, y = self._vox[list(self._xy_idx[axis])] self._images['cursor_v'][axis].set_xdata([x, x]) self._images['cursor_h'][axis].set_ydata([y, y]) self._zoom(0) # doesn't actually zoom just resets view to center self._update_images(draw=True) self._update_moved() def _show_help(self): """Show the help menu.""" QMessageBox.information( self, 'Help', "Help:\n" "'+'/'-': zoom\nleft/right arrow: left/right\n" "up/down arrow: superior/inferior\n" "left angle bracket/right angle bracket: anterior/posterior") def _key_press_event(self, event): """Execute functions when the user presses a key.""" if event.key() == 'escape': self.close() if event.text() == 'h': self._show_help() if event.text() in ('=', '+', '-'): self._zoom(sign=-2 * (event.text() == '-') + 1, draw=True) # Changing slices if event.key() in (QtCore.Qt.Key_Up, QtCore.Qt.Key_Down, QtCore.Qt.Key_Left, QtCore.Qt.Key_Right, QtCore.Qt.Key_Comma, QtCore.Qt.Key_Period, QtCore.Qt.Key_PageUp, QtCore.Qt.Key_PageDown): ras = np.array(self._ras) if event.key() in (QtCore.Qt.Key_Up, QtCore.Qt.Key_Down): ras[2] += 2 * (event.key() == QtCore.Qt.Key_Up) - 1 elif event.key() in (QtCore.Qt.Key_Left, QtCore.Qt.Key_Right): ras[0] += 2 * (event.key() == QtCore.Qt.Key_Right) - 1 else: ras[1] += 2 * (event.key() == QtCore.Qt.Key_PageUp or event.key() == QtCore.Qt.Key_Period) - 1 self._set_ras(ras) def _on_click(self, event, axis): """Move to view on MRI and CT on click.""" if event.inaxes is self._figs[axis].axes[0]: # Data coordinates are voxel coordinates pos = (event.xdata, event.ydata) logger.info(f'Clicked {"XYZ"[axis]} ({axis}) axis at pos {pos}') xyz = self._vox xyz[list(self._xy_idx[axis])] = pos logger.debug(f'Using voxel {list(xyz)}') ras = apply_trans(self._vox_ras_t, xyz) self._set_ras(ras) def _update_moved(self): """Update when cursor position changes.""" self._RAS_textbox.setPlainText('{:.2f}, {:.2f}, {:.2f}'.format( *self._ras)) self._VOX_textbox.setPlainText('{:3d}, {:3d}, {:3d}'.format( *self._current_slice)) self._intensity_label.setText('intensity = {:.2f}'.format( self._base_data[tuple(self._current_slice)])) @safe_event def closeEvent(self, event): """Clean up upon closing the window.""" self._renderer.plotter.close() self.close()