diff --git a/doc/source/pyplots/annotation.py b/doc/source/pyplots/annotation.py index b1f573d..5030281 100644 --- a/doc/source/pyplots/annotation.py +++ b/doc/source/pyplots/annotation.py @@ -11,6 +11,6 @@ annotation[Segment(6, 8)] = 'Bob' annotation[Segment(12, 18)] = 'Carol' annotation[Segment(7, 20)] = 'Alice' -notebook.plot_annotation(annotation, legend=True, time=True) +notebook.plot_annotation(annotation, legend=True, time=True, arrangement="pack") plt.show() diff --git a/doc/source/pyplots/introduction.py b/doc/source/pyplots/introduction.py index 54a6ecd..5e7032e 100644 --- a/doc/source/pyplots/introduction.py +++ b/doc/source/pyplots/introduction.py @@ -12,7 +12,7 @@ annotation[Segment(6, 8)] = 'Bob' annotation[Segment(12, 18)] = 'Carol' annotation[Segment(7, 20)] = 'Alice' -notebook.plot_annotation(annotation, legend=True, time=False) +notebook.plot_annotation(annotation, legend=True, time=False, arrangement="pack") # plot timeline plt.subplot(212) diff --git a/pyannote/core/annotation.py b/pyannote/core/annotation.py index 19e4311..9144ff3 100755 --- a/pyannote/core/annotation.py +++ b/pyannote/core/annotation.py @@ -1430,4 +1430,4 @@ def _repr_png_(self): return None from .notebook import repr_annotation - return repr_annotation(self) + return repr_annotation(self, arrangement="pack") diff --git a/pyannote/core/notebook.py b/pyannote/core/notebook.py index a08b8e1..82affea 100644 --- a/pyannote/core/notebook.py +++ b/pyannote/core/notebook.py @@ -261,7 +261,7 @@ def __call__(self, resource: Resource, self.plot_timeline(resource, time=time) elif isinstance(resource, Annotation): - self.plot_annotation(resource, time=time, legend=legend) + self.plot_annotation(resource, time=time, legend=legend, arrangement="pack") elif isinstance(resource, SlidingWindowFeature): self.plot_feature(resource, time=time) @@ -288,7 +288,20 @@ def plot_timeline(self, timeline: Timeline, ax=None, time=True): # ax.set_aspect(3. / self.crop.duration) - def plot_annotation(self, annotation: Annotation, ax=None, time=True, legend=True): + def plot_annotation(self, annotation: Annotation, ax=None, time=True, legend=True, arrangement="pack"): + """ plots annotation + + Parameters + ---------- + arrangement : str + To plot segments, they are grouped by certain rules, and each group is plotted + on a distinct line. `stack` will group segments according to their label. + `pack` will group segments optimally by timestamp, regardless of label. + + Returns + ------- + + """ if not self.crop: self.crop = annotation.get_timeline(copy=False).extent() @@ -299,10 +312,18 @@ def plot_annotation(self, annotation: Annotation, ax=None, time=True, legend=Tru ax = self.setup(ax=ax, time=time) - for (segment, track, label), y in zip( - cropped.itertracks(yield_label=True), - self.get_y(segments)): - self.draw_segment(ax, segment, y, label=label) + if arrangement == "pack": + for (segment, track, label), y in zip( + cropped.itertracks(yield_label=True), + self.get_y(segments)): + self.draw_segment(ax, segment, y, label=label) + + elif arrangement == "stack": + labels_dict = {label: i for i, label in enumerate(set(labels))} + for (segment, track, label) in \ + cropped.itertracks(yield_label=True): + y = 1.0 - 1.0 / (len(labels) + 1) * (1 + labels_dict.get(label)) + self.draw_segment(ax, segment, y, label=label) if legend: H, L = ax.get_legend_handles_labels() @@ -375,13 +396,13 @@ def repr_timeline(timeline: Timeline): return data -def repr_annotation(annotation: Annotation): +def repr_annotation(annotation: Annotation, arrangement="pack"): """Get `png` data for `annotation`""" import matplotlib.pyplot as plt figsize = plt.rcParams['figure.figsize'] plt.rcParams['figure.figsize'] = (notebook.width, 2) fig, ax = plt.subplots() - notebook.plot_annotation(annotation, ax=ax) + notebook.plot_annotation(annotation, ax=ax, arrangement=arrangement) data = print_figure(fig, 'png') plt.close(fig) plt.rcParams['figure.figsize'] = figsize