Skip to content

Commit d6c6dc3

Browse files
jaeileppagramfort
authored andcommitted
[MRG] Interactive annotations. (mne-tools#3821)
* Interactive annotations.
1 parent 7e0b06d commit d6c6dc3

8 files changed

Lines changed: 533 additions & 67 deletions

File tree

doc/whats_new.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,8 @@ Changelog
7171

7272
- Add support for NeuroScan files with event type 3 in :func:`mne.io.read_raw_cnt` by `Marijn van Vliet`_
7373

74+
- Add interactive annotation mode to :meth:`mne.io.Raw.plot` (accessed by pressing 'a') by `Jaakko Leppakangas`_
75+
7476
BUG
7577
~~~
7678

mne/annotations.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,36 @@ def __init__(self, onset, duration, description,
8484
self.duration = duration
8585
self.description = np.array(description, dtype=str)
8686

87+
def append(self, onset, duration, description):
88+
"""Add an annotated segment. Operates inplace.
89+
90+
Parameters
91+
----------
92+
onset : float
93+
Annotation time onset from the beginning of the recording in
94+
seconds.
95+
duration : float
96+
Duration of the annotation in seconds.
97+
description : str
98+
Description for the annotation. To reject epochs, use description
99+
starting with keyword 'bad'
100+
"""
101+
self.onset = np.append(self.onset, onset)
102+
self.duration = np.append(self.duration, duration)
103+
self.description = np.append(self.description, description)
104+
105+
def delete(self, idx):
106+
"""Remove an annotation. Operates inplace.
107+
108+
Parameters
109+
----------
110+
idx : int | list of int
111+
Index of the annotation to remove.
112+
"""
113+
self.onset = np.delete(self.onset, idx)
114+
self.duration = np.delete(self.duration, idx)
115+
self.description = np.delete(self.description, idx)
116+
87117

88118
def _combine_annotations(annotations, last_samps, first_samps, sfreq):
89119
"""Helper for combining a tuple of annotations."""

mne/tests/test_annotations.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,12 @@ def test_annotations():
6262
raws.append(raw)
6363
raw = concatenate_raws(raws)
6464
assert_array_equal(raw.annotations.onset, [1., 2., 11., 12., 21., 22.])
65-
65+
raw.annotations.delete(2)
66+
assert_array_equal(raw.annotations.onset, [1., 2., 12., 21., 22.])
67+
raw.annotations.append(5, 1.5, 'y')
68+
assert_array_equal(raw.annotations.onset, [1., 2., 12., 21., 22., 5])
69+
assert_array_equal(raw.annotations.duration, [.5, .5, .5, .5, .5, 1.5])
70+
assert_array_equal(raw.annotations.description, ['x', 'x', 'x', 'x', 'x',
71+
'y'])
6672

6773
run_tests_if_main()

mne/viz/raw.py

Lines changed: 16 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,8 @@
2626
_helper_raw_resize, _select_bads, _onclick_help,
2727
_setup_browser_offsets, _compute_scalings, plot_sensors,
2828
_radio_clicked, _set_radio_button, _handle_topomap_bads,
29-
_change_channel_group)
29+
_change_channel_group, _plot_annotations)
3030
from ..defaults import _handle_default
31-
from ..annotations import _onset_to_seconds
3231
from .evoked import _plot_lines
3332

3433

@@ -81,6 +80,8 @@ def _update_raw_data(params):
8180
def _pick_bad_channels(event, params):
8281
"""Selecting / drop bad channels onpick."""
8382
# Both bad lists are updated. params['info'] used for colors.
83+
if params['fig_annotation'] is not None:
84+
return
8485
bads = params['raw'].info['bads']
8586
params['info']['bads'] = _select_bads(event, params, bads)
8687
_plot_update_raw_proj(params, None)
@@ -310,7 +311,7 @@ def plot_raw(raw, events=None, duration=10.0, start=0.0, n_channels=20,
310311
n_channels=n_channels, scalings=scalings, types=types,
311312
n_times=n_times, event_times=event_times, inds=inds,
312313
event_nums=event_nums, clipping=clipping, fig_proj=None,
313-
first_time=first_time)
314+
first_time=first_time, added_label=list())
314315

315316
if isinstance(order, string_types) and order in ['selection', 'position']:
316317
params['fig_selection'] = fig_selection
@@ -331,31 +332,7 @@ def plot_raw(raw, events=None, duration=10.0, start=0.0, n_channels=20,
331332
bad_color=bad_color, event_lines=event_lines,
332333
event_color=event_color)
333334

334-
if raw.annotations is not None:
335-
segments = list()
336-
segment_colors = dict()
337-
# sort the segments by start time
338-
ann_order = raw.annotations.onset.argsort(axis=0)
339-
descriptions = raw.annotations.description[ann_order]
340-
color_keys = set(descriptions)
341-
color_vals = np.linspace(0, 1, len(color_keys))
342-
for idx, key in enumerate(color_keys):
343-
if key.lower().startswith('bad'):
344-
segment_colors[key] = 'red'
345-
else:
346-
segment_colors[key] = plt.cm.summer(color_vals[idx])
347-
params['segment_colors'] = segment_colors
348-
for idx, onset in enumerate(raw.annotations.onset[ann_order]):
349-
annot_start = _onset_to_seconds(raw, onset) + first_time
350-
annot_end = annot_start + raw.annotations.duration[ann_order][idx]
351-
segments.append([annot_start, annot_end])
352-
ylim = params['ax_hscroll'].get_ylim()
353-
dscr = descriptions[idx]
354-
params['ax_hscroll'].fill_betweenx(ylim, annot_start, annot_end,
355-
alpha=0.3,
356-
color=segment_colors[dscr])
357-
params['segments'] = np.array(segments)
358-
params['annot_description'] = descriptions
335+
_plot_annotations(raw, params)
359336

360337
params['update_fun'] = partial(_update_raw_data, params=params)
361338
params['pick_bads_fun'] = partial(_pick_bad_channels, params=params)
@@ -396,13 +373,14 @@ def plot_raw(raw, events=None, duration=10.0, start=0.0, n_channels=20,
396373
# deal with projectors
397374
if show_options is True:
398375
_toggle_options(None, params)
376+
377+
callback_close = partial(_close_event, params=params)
378+
params['fig'].canvas.mpl_connect('close_event', callback_close)
399379
# initialize the first selection set
400380
if isinstance(order, string_types) and order in ['selection', 'position']:
401381
_radio_clicked(fig_selection.radio.labels[0]._text, params)
402382
callback_selection_key = partial(_selection_key_press, params=params)
403383
callback_selection_scroll = partial(_selection_scroll, params=params)
404-
callback_close = partial(_close_event, params=params)
405-
params['fig'].canvas.mpl_connect('close_event', callback_close)
406384
params['fig_selection'].canvas.mpl_connect('close_event',
407385
callback_close)
408386
params['fig_selection'].canvas.mpl_connect('key_press_event',
@@ -439,7 +417,10 @@ def _selection_key_press(event, params):
439417
def _close_event(event, params):
440418
"""Callback for closing of raw browser with selections."""
441419
import matplotlib.pyplot as plt
442-
plt.close(params['fig_selection'])
420+
if 'fig_selection' in params:
421+
plt.close(params['fig_selection'])
422+
if 'fig_annotation' is not None:
423+
plt.close(params['fig_annotation'])
443424
plt.close(params['fig'])
444425

445426

@@ -785,6 +766,8 @@ def _prepare_mne_browse_raw(params, title, bgcolor, color, bad_color, inds,
785766
params['lines'] = [ax.plot([np.nan], antialiased=False, linewidth=0.5)[0]
786767
for _ in range(n_ch)]
787768
ax.set_yticklabels(['X' * max([len(ch) for ch in info['ch_names']])])
769+
params['fig_annotation'] = None
770+
params['segment_line'] = None
788771

789772

790773
def _plot_raw_traces(params, color, bad_color, event_lines=None,
@@ -874,8 +857,8 @@ def _plot_raw_traces(params, color, bad_color, event_lines=None,
874857
break # Since the segments are sorted by t_start
875858
if segment[1] < times[0] + params['first_time']:
876859
continue
877-
start = segment[0]
878-
end = segment[1]
860+
start = max(segment[0], times[0] + params['first_time'])
861+
end = min(times[-1] + params['first_time'], segment[1])
879862
dscr = params['annot_description'][idx]
880863
segment_color = params['segment_colors'][dscr]
881864
params['ax'].fill_betweenx(ylim, start, end, color=segment_color,

mne/viz/tests/test_raw.py

Lines changed: 61 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from mne import read_events, pick_types, Annotations
1212
from mne.io import read_raw_fif
1313
from mne.utils import requires_version, run_tests_if_main
14-
from mne.viz.utils import _fake_click
14+
from mne.viz.utils import _fake_click, _annotation_radio_clicked
1515
from mne.viz import plot_raw, plot_sensors
1616

1717
# Set our plotters to test mode
@@ -41,7 +41,6 @@ def _get_events():
4141
return read_events(event_name)
4242

4343

44-
@requires_version('matplotlib', '1.2')
4544
def test_plot_raw():
4645
"""Test plotting of raw data."""
4746
import matplotlib.pyplot as plt
@@ -53,7 +52,8 @@ def test_plot_raw():
5352
# test mouse clicks
5453
x = fig.get_axes()[0].lines[1].get_xdata().mean()
5554
y = fig.get_axes()[0].lines[1].get_ydata().mean()
56-
data_ax = fig.get_axes()[0]
55+
data_ax = fig.axes[0]
56+
5757
_fake_click(fig, data_ax, [x, y], xform='data') # mark a bad channel
5858
_fake_click(fig, data_ax, [x, y], xform='data') # unmark a bad channel
5959
_fake_click(fig, data_ax, [0.5, 0.999]) # click elsewhere in 1st axes
@@ -75,22 +75,9 @@ def test_plot_raw():
7575
# _fake_click(ssp_fig, ssp_fig.get_axes()[0], pos, xform='data') # off
7676
# _fake_click(ssp_fig, ssp_fig.get_axes()[0], pos, xform='data') # on
7777
# test keypresses
78-
fig.canvas.key_press_event('escape')
79-
fig.canvas.key_press_event('down')
80-
fig.canvas.key_press_event('up')
81-
fig.canvas.key_press_event('right')
82-
fig.canvas.key_press_event('left')
83-
fig.canvas.key_press_event('o')
84-
fig.canvas.key_press_event('-')
85-
fig.canvas.key_press_event('+')
86-
fig.canvas.key_press_event('=')
87-
fig.canvas.key_press_event('pageup')
88-
fig.canvas.key_press_event('pagedown')
89-
fig.canvas.key_press_event('home')
90-
fig.canvas.key_press_event('end')
91-
fig.canvas.key_press_event('?')
92-
fig.canvas.key_press_event('f11')
93-
fig.canvas.key_press_event('escape')
78+
for key in ['down', 'up', 'right', 'left', 'o', '-', '+', '=',
79+
'pageup', 'pagedown', 'home', 'end', '?', 'f11', 'escape']:
80+
fig.canvas.key_press_event(key)
9481
# Color setting
9582
assert_raises(KeyError, raw.plot, event_color={0: 'r'})
9683
assert_raises(TypeError, raw.plot, event_color={'foo': 'r'})
@@ -134,6 +121,61 @@ def test_plot_raw():
134121
plt.close('all')
135122

136123

124+
def test_plot_annotations():
125+
"""Test annotation mode of the plotter."""
126+
import matplotlib.pyplot as plt
127+
raw = _get_raw()
128+
fig = raw.plot()
129+
data_ax = fig.axes[0]
130+
fig.canvas.key_press_event('a') # annotation mode
131+
# modify description
132+
ann_fig = plt.gcf()
133+
for key in ' test':
134+
ann_fig.canvas.key_press_event(key)
135+
ann_fig.canvas.key_press_event('enter')
136+
137+
ann_fig = plt.gcf()
138+
# XXX: _fake_click raises an error on Agg backend
139+
_annotation_radio_clicked('', ann_fig.radio, data_ax.selector)
140+
141+
# draw annotation
142+
_fake_click(fig, data_ax, [1., 1.], xform='data', button=1, kind='press')
143+
_fake_click(fig, data_ax, [5., 1.], xform='data', button=1, kind='motion')
144+
_fake_click(fig, data_ax, [5., 1.], xform='data', button=1, kind='release')
145+
# hover event
146+
_fake_click(fig, data_ax, [4.5, 1.], xform='data', button=None,
147+
kind='motion')
148+
_fake_click(fig, data_ax, [4.7, 1.], xform='data', button=None,
149+
kind='motion')
150+
# modify annotation from end
151+
_fake_click(fig, data_ax, [5., 1.], xform='data', button=1, kind='press')
152+
_fake_click(fig, data_ax, [2.5, 1.], xform='data', button=1, kind='motion')
153+
_fake_click(fig, data_ax, [2.5, 1.], xform='data', button=1,
154+
kind='release')
155+
# modify annotation from beginning
156+
_fake_click(fig, data_ax, [1., 1.], xform='data', button=1, kind='press')
157+
_fake_click(fig, data_ax, [1.1, 1.], xform='data', button=1, kind='motion')
158+
_fake_click(fig, data_ax, [1.1, 1.], xform='data', button=1,
159+
kind='release')
160+
assert_equal(len(raw.annotations.onset), 1)
161+
assert_equal(len(raw.annotations.duration), 1)
162+
assert_equal(len(raw.annotations.description), 1)
163+
assert_equal(raw.annotations.description[0], 'BAD test')
164+
165+
# draw another annotation merging the two
166+
_fake_click(fig, data_ax, [5.5, 1.], xform='data', button=1, kind='press')
167+
_fake_click(fig, data_ax, [2., 1.], xform='data', button=1, kind='motion')
168+
_fake_click(fig, data_ax, [2., 1.], xform='data', button=1, kind='release')
169+
# delete the annotation
170+
_fake_click(fig, data_ax, [1.5, 1.], xform='data', button=3, kind='press')
171+
fig.canvas.key_press_event('a') # exit annotation mode
172+
plt.close('all')
173+
174+
assert_equal(len(raw.annotations.onset), 0)
175+
assert_equal(len(raw.annotations.duration), 0)
176+
assert_equal(len(raw.annotations.description), 0)
177+
178+
137179
@requires_version('scipy', '0.10')
138180
def test_plot_raw_filtered():
139181
"""Test filtering of raw plots."""
@@ -187,7 +229,6 @@ def test_plot_raw_psd():
187229
assert_raises(ValueError, raw.plot_psd)
188230

189231

190-
@requires_version('matplotlib', '1.2')
191232
def test_plot_sensors():
192233
"""Test plotting of sensor array."""
193234
import matplotlib.pyplot as plt

0 commit comments

Comments
 (0)