Update src/demo/utils.py

src/demo/utils.py

@@ -44,12 +44,16 @@ def process_audio(model, audio, config):
 
 
 def get_spec_pil(model, audio, config):
-    fbank, spec_plot = process_audio(model, audio, config)
-    buf = io.BytesIO()
-    spec_plot.figure.savefig(buf, format='png')
-    buf.seek(0)
-    pil_spec = Image.open(buf)
-    plt.close()
+    try:
+        fbank, spec_plot = process_audio(model, audio, config)
+        buf = io.BytesIO()
+        spec_plot.figure.savefig(buf, format='png')
+        buf.seek(0)
+        pil_spec = Image.open(buf)
+        plt.close()
+    except:
+        print("Warning: the streaming is not ready. Please repeate uploading again.")
+        fbank, pil_spec = None, None
     return fbank, pil_spec
 
 
@@ -74,7 +78,8 @@ def get_mask_region(img):
     # Use the channel of opacity as mask
     mask = pil_to_tensor(layers[0])[-1,:,:]  # RGBA
     mask = mask.permute(1, 0)  # (F, T) -> (T, F)
-
+    # Flip the freq axis to ensure the orignal point on the top left
+    mask = mask.flip(1)
     mask = (mask > 0).float()
 
     # Rescale mask to spectrum size
@@ -85,6 +90,8 @@ def get_mask_region(img):
 def get_mask_regions(img):
     def _prepare_mask(m):
         m = m.permute(1, 0)
+        # Flip the freq axis to ensure the orignal point on the top left
+        m = m.flip(1)
         m = (m > 0).float()
         m = F.interpolate(m.unsqueeze(0).unsqueeze(0), SPEC_RES).squeeze()
         return m
@@ -112,7 +119,8 @@ def update_reference_spec(ref_spec_pil_ori, mask_src, dt, df, resize_scale_t, re
     if mask_src is not None:
         mask_ref = get_edit_mask(
             mask_src, dx=df, dy=dt, 
-            resize_scale_x=resize_scale_f, resize_scale_y=resize_scale_t,
+            resize_scale_x=resize_scale_f, 
+            resize_scale_y=resize_scale_t,
             )
         mask_ref = mask_ref.float()  # match the PIL format, channel last
         mask_ref_pil = F.interpolate(mask_ref.unsqueeze(0).unsqueeze(0), DESPLAY_RES).squeeze()
@@ -121,6 +129,9 @@ def update_reference_spec(ref_spec_pil_ori, mask_src, dt, df, resize_scale_t, re
         if mask_ref_pil.ndim > 2:
             mask_ref_pil = mask_ref_pil.squeeze()
         mask_ref_pil = mask_ref_pil.permute(1, 0)
+        # De-flip freq exis to match pil imshow style
+        mask_ref_pil = mask_ref_pil.flip(0)
+        mask_ref_pil = mask_ref_pil * 0.5  # for transparency
 
         # Convert to PIL
         mask_ref_pil = to_pil_image(mask_ref_pil).convert("L")