Merge branch 'main' of https://huggingface.co/spaces/anhnv125/FRN

README.md

@@ -3,8 +3,10 @@ title: FRN
 emoji: 📉
 colorFrom: gray
 colorTo: red
-sdk: static
-pinned: false
+sdk: streamlit
+pinned: true
+app_file: app.py
+sdk_version: 1.10.0
 ---
 
 # FRN - Full-band Recurrent Network Official Implementation

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import librosa
+import soundfile as sf
 import librosa.display
 from config import CONFIG
 import torch
@@ -9,7 +10,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
 
-@st.cache_resource
+@st.cache
 def load_model():
     path = 'lightning_logs/version_0/checkpoints/frn.onnx'
     onnx_model = onnx.load(path)
@@ -87,9 +88,9 @@ target = target[:packet_size * (len(target) // packet_size)]
 st.subheader('Original audio')
 st.audio(uploaded_file)
 
-st.subheader('Choose loss packet percentage')
-loss_percent = st.radio('Loss percentage', ['10%', '20%', '30%', '40%'])
-loss_percent = float(loss_percent[:-1])/100
+st.subheader('Choose expected packet loss rate')
+slider = [st.slider("Expected loss rate for Markov Chain loss generator", 0, 100, step=1)]
+loss_percent = float(slider[0])/100
 mask_gen = MaskGenerator(is_train=False, probs=[(1 - loss_percent, loss_percent)])
 lossy_input = target.copy().reshape(-1, packet_size)
 mask = mask_gen.gen_mask(len(lossy_input), seed=0)[:, np.newaxis]
@@ -109,9 +110,12 @@ if st.button('Conceal lossy audio!'):
         fig = visualize(target, lossy_input, output)
         st.pyplot(fig)
     st.success('Done!')
+    sf.write('target.wav', target, sr)
+    sf.write('lossy.wav', lossy_input, sr)
+    sf.write('enhanced.wav', output, sr)
     st.text('Original audio')
-    st.audio(target, sample_rate=sr)
+    st.audio('target.wav')
     st.text('Lossy audio')
-    st.audio(lossy_input, sample_rate=sr)
+    st.audio('lossy.wav')
     st.text('Enhanced audio')
-    st.audio(output, sample_rate=sr)
+    st.audio('enhanced.wav')

requirements.txt

@@ -12,6 +12,6 @@ soundfile==0.11.0
 torch==1.13.1
 torchmetrics==0.11.0
 tqdm==4.64.0
-stoi==0.3.3
+pystoi==0.3.3
 pesq==0.0.4
 onnx==1.13.0