From 1b700124b5bfe642a3973abba3b21731b78fd9cd Mon Sep 17 00:00:00 2001
From: Lucas Gautheron <lucas.gautheron@gmail.com>
Date: Thu, 4 Mar 2021 13:42:02 +0100
Subject: [PATCH] annotations visualization

---
 ChildProject/annotations.py | 28 +++++++++++++++++++++++++++-
 ChildProject/cmdline.py     | 25 +++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/ChildProject/annotations.py b/ChildProject/annotations.py
index 20a19bd03..737ec24a5 100644
--- a/ChildProject/annotations.py
+++ b/ChildProject/annotations.py
@@ -982,4 +982,30 @@ def get_vc_stats(self, segments: pd.DataFrame, turntakingthresh: float = 1):
             voc_count = ('duration', 'count'),
             turns = ('turn', 'sum'),
             cds_dur = ('cds', 'sum')
-        )
\ No newline at end of file
+        )
+
+    def plot(self, annotations):
+        from matplotlib import pyplot as plt
+        from pyannote.core import Timeline, Annotation, Segment, notebook
+
+        nrows = len(annotations['recording_filename'].unique())
+        fig, ax = plt.subplots(nrows = nrows, ncols = 1)
+        fig.set_figwidth(20)
+        fig.set_figheight(nrows * 2)
+
+        i = 0
+        for recording_filename, recording_annotations in annotations.groupby('recording_filename'):
+            pyannotation = Annotation()
+
+            for annotation in recording_annotations.to_dict(orient = 'records'):
+                start = annotation['range_onset'] + annotation['time_seek']
+                end = annotation['range_offset'] + annotation['time_seek']
+
+                pyannotation[Segment(start, end), annotation['set']] = annotation['set']
+
+                print(annotation)
+
+            notebook.plot_annotation(pyannotation, ax = ax[i], legend = True, time = True)
+            i += 1
+        
+        return fig, ax
diff --git a/ChildProject/cmdline.py b/ChildProject/cmdline.py
index 5958fc66f..cd6e992c5 100755
--- a/ChildProject/cmdline.py
+++ b/ChildProject/cmdline.py
@@ -170,6 +170,31 @@ def rename_annotations(args):
     am.read()
     am.rename_set(args.set, args.new_set, recursive = args.recursive, ignore_errors = args.ignore_errors)
 
+@subcommand([
+    arg("source", help = "project path"),
+    arg("--recordings", help = 'list of recordings to plot', nargs = '*', default = [])
+])
+def plot_annotations(args):
+    from matplotlib import pyplot as plt
+    """show a diagram representing available annotations for each recording"""
+
+    project = ChildProject(args.source)
+    errors, warnings = project.validate(ignore_files = True)
+
+    if len(errors) > 0:
+        print("validation failed, {} error(s) occured".format(len(errors)), file = sys.stderr)
+        sys.exit(1)
+
+    am = AnnotationManager(project)
+    am.read()
+    annotations = am.annotations
+
+    if len(args.recordings):
+        annotations = annotations[annotations['recording_filename'].isin(args.recordings)]
+
+    fig, ax = am.plot(annotations)
+    plt.show()
+
 @subcommand([
     arg("dataset", help = "dataset to install. Should be a valid repository name at https://github.com/LAAC-LSCP. (e.g.: solomon-data)"),
     arg("--destination", help = "destination path", required = False, default = ""),