From ffe51515fcc5916221f58664de1dc4fe56ce7cb8 Mon Sep 17 00:00:00 2001
From: koukyo1994 <arabiannight1994@yahoo.co.jp>
Date: Fri, 26 Jun 2020 23:56:28 +0900
Subject: [PATCH] add option to show annotation

---
 components/__init__.py |  2 +-
 components/plots.py    | 67 ++++++++++++++++++++++++++++++++++++++++++
 main.py                | 27 ++++++++++++++---
 utils/__init__.py      |  2 +-
 utils/io.py            |  7 +++++
 5 files changed, 99 insertions(+), 6 deletions(-)

diff --git a/components/__init__.py b/components/__init__.py
index 8ce917d..19f570e 100644
--- a/components/__init__.py
+++ b/components/__init__.py
@@ -1,3 +1,3 @@
 from .base import *
-from .plots import waveplot, specshow
+from .plots import waveplot, specshow, waveplot_with_annotation
 from .preprocessing import preprocess_on_wave
diff --git a/components/plots.py b/components/plots.py
index 2ca2941..e7184d5 100644
--- a/components/plots.py
+++ b/components/plots.py
@@ -2,6 +2,7 @@
 import librosa.display as display
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import streamlit as st
 
 
@@ -42,6 +43,72 @@ def waveplot(y: np.ndarray, sr: int, processed=None):
         st.pyplot()
 
 
+def waveplot_with_annotation(y: np.ndarray,
+                             sr: int,
+                             annotation: pd.DataFrame,
+                             filename: str,
+                             processed=None):
+    plot_wave = st.checkbox("Waveplot")
+    if filename.endswith(".mp3"):
+        filename = filename.replace(".mp3", ".wav")
+    events = annotation.query(f"filename == '{filename}'")
+    colors = [
+        "#bf6565", "#ac7ceb", "#e3e176", "#f081e1", "#e8cb6b", "#25b4db",
+        "#fa787e", "#a9f274", "#1d7335", "#797fb3"
+    ]
+    if plot_wave:
+        st.sidebar.markdown("#### Waveplot settings")
+        start_second = st.sidebar.number_input(
+            "start second",
+            min_value=0,
+            max_value=len(y) // sr,
+            value=0,
+            step=1,
+            key="waveplot_start")
+        end_second = st.sidebar.number_input(
+            "end second",
+            min_value=0,
+            max_value=len(y) // sr,
+            value=len(y) // sr,
+            step=1,
+            key="waveplot_end")
+
+        start_index = start_second * sr
+        if end_second == len(y) // sr:
+            end_index = len(y)
+        else:
+            end_index = end_second * sr
+        events_in_period = events.query(
+            f"onset > {start_second} & offset < {end_second}")
+        uniq_labels = events_in_period["ebird_code"].unique().tolist()
+        plt.figure(figsize=(12, 4))
+        plt.grid(True)
+        display.waveplot(y[start_index:end_index], sr=sr, alpha=0.5)
+
+        used_color = []  # type: ignore
+        for i, event in events_in_period.iterrows():
+            onset = event.onset
+            offset = event.offset
+            color = colors[uniq_labels.index(event.ebird_code)]
+            if color not in used_color:
+                label = event.ebird_code
+                used_color.append(color)
+            else:
+                label = "_" + event.ebird_code
+            plt.axvspan(onset, offset, facecolor=color, alpha=0.5, label=label)
+
+        plt.legend()
+
+        if processed is not None:
+            display.waveplot(
+                processed[start_index:end_index],
+                sr=sr,
+                alpha=0.5,
+                color="red")
+
+        st.pyplot()
+
+
 @st.cache
 def melspectrogram(y: np.ndarray, params: dict, log=True):
     melspec = librosa.feature.melspectrogram(y=y, **params)
diff --git a/main.py b/main.py
index 41cc9e0..56bdc65 100644
--- a/main.py
+++ b/main.py
@@ -29,14 +29,33 @@
             options=["normal", "preprocessing", "augmentations"])
         utils.display_media_audio(audio_path, second)
 
+        annotation = st.sidebar.file_uploader(
+            "Upload annotation file if exist")
+        if annotation is not None:
+            event_level_annotation = utils.read_csv(annotation)
+        else:
+            event_level_annotation = None
+
         y = utils.read_audio(audio_path, audio_info, sr=sr)
         if options == "preprocessing":
             y_processed = C.preprocess_on_wave(y, sr=sr)
             if y_processed is not None:
                 st.text("Processed audio")
                 utils.display_media_audio_from_ndarray(y_processed, sr)
-                C.waveplot(y, sr, y_processed)
-                C.specshow(y, sr, y_processed)
+                if event_level_annotation is None:
+                    C.waveplot(y, sr, y_processed)
+                    C.specshow(y, sr, y_processed)
+                else:
+                    C.waveplot_with_annotation(y, sr, event_level_annotation,
+                                               audio_file_name, y_processed)
         else:
-            C.waveplot(y, sr)
-            C.specshow(y, sr)
+            if event_level_annotation is None:
+                C.waveplot(y, sr)
+                C.specshow(y, sr)
+            else:
+                C.waveplot_with_annotation(
+                    y,
+                    sr,
+                    event_level_annotation,
+                    audio_file_name,
+                    processed=None)
diff --git a/utils/__init__.py b/utils/__init__.py
index c4a604b..234a8ae 100644
--- a/utils/__init__.py
+++ b/utils/__init__.py
@@ -1 +1 @@
-from .io import check_folder, check_audio_info, display_media_audio, read_audio, display_media_audio_from_ndarray
+from .io import check_folder, check_audio_info, display_media_audio, read_audio, display_media_audio_from_ndarray, read_csv
diff --git a/utils/io.py b/utils/io.py
index d76d439..aec85db 100644
--- a/utils/io.py
+++ b/utils/io.py
@@ -6,12 +6,19 @@
 
 import librosa
 import numpy as np
+import pandas as pd
 import streamlit as st
 
 from pathlib import Path
 from typing import Optional
 
 
+@st.cache
+def read_csv(uploaded_file):
+    df = pd.read_csv(uploaded_file)
+    return df
+
+
 @st.cache
 def read_audio(path: Path, info: dict, sr: Optional[int] = None):
     if sr is None: