From 995d0e6624777c424b13819638e789c9aa90723e Mon Sep 17 00:00:00 2001
From: lhzstar <2792382045@qq.com>
Date: Mon, 26 Jun 2023 14:28:43 +0800
Subject: [PATCH] add progress bar for UMAP

---
 toolbox/__init__.py | 6 ------
 toolbox/ui.py       | 6 ++++++
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/toolbox/__init__.py b/toolbox/__init__.py
index 2fa95b6..326a756 100644
--- a/toolbox/__init__.py
+++ b/toolbox/__init__.py
@@ -208,11 +208,8 @@ class Toolbox:
         self.ui.register_utterance(utterance)
 
         # Plot it
-        self.ui.log("Generating UMAP...")
-        self.ui.set_loading(1)
         self.ui.draw_embed(embed, name, "current")
         self.ui.draw_umap_projections(self.utterances)
-        self.ui.set_loading(0)
 
     def clear_utterances(self):
         self.utterances.clear()
@@ -356,11 +353,8 @@ class Toolbox:
         self.utterances.add(utterance)
 
         # Plot it
-        self.ui.log("Generating UMAP...")
-        self.ui.set_loading(1)
         self.ui.draw_embed(embed, name, "generated")
         self.ui.draw_umap_projections(self.utterances)
-        self.ui.set_loading(0)
 
     def init_encoder(self):
         model_fpath = self.ui.current_encoder_fpath
diff --git a/toolbox/ui.py b/toolbox/ui.py
index 5ba244c..c8fb087 100644
--- a/toolbox/ui.py
+++ b/toolbox/ui.py
@@ -86,6 +86,8 @@ class UI(QDialog):
             self.vocode_button.setDisabled(spec is None)
 
     def draw_umap_projections(self, utterances: Set[Utterance]):
+        def umap_progress(i, seq_len):
+            self.set_loading(i, seq_len)
         self.umap_ax.clear()
 
         speakers = np.unique([u.speaker_name for u in utterances])
@@ -110,14 +112,18 @@ class UI(QDialog):
             projections = reducer.fit_transform(embeds)
 
             speakers_done = set()
+            i = 0
             for projection, utterance in zip(projections, utterances):
+                i+=1
                 color = colors[utterance.speaker_name]
                 mark = "x" if "_gen_" in utterance.name else "o"
                 label = None if utterance.speaker_name in speakers_done else utterance.speaker_name
                 speakers_done.add(utterance.speaker_name)
                 self.umap_ax.scatter(projection[0], projection[1], c=[color], marker=mark,
                                      label=label)
+                self.set_loading(i, projections.shape[0])
             self.umap_ax.legend(prop={'size': 10})
+            self.set_loading(0)
 
         # Draw the plot
         self.umap_ax.set_aspect("equal", "datalim")