Update app.py

app.py

@@ -1,10 +1,11 @@
 import os
 import torch
+import torch.nn.functional as F
 import gradio as gr
 from huggingface_hub import InferenceClient
 from model import predict_params, AudioDataset
 import torchaudio
-# TODO: Que no diga lo de que no hay 1s_normal al predecir
+
 token = os.getenv("HF_TOKEN")
 client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct", token=token)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -44,7 +45,7 @@ def predict(audio_path_pred):
 def predict_stream(audio_path_stream):
     with torch.no_grad():
         logits = call(audio_path_stream, model=model_mon, dataset_path="data/baby_cry_detection", filter_white_noise=False, undersample_normal=False)
-        probabilities = torch.nn.functional.softmax(logits, dim=-1)
+        probabilities = F.softmax(logits, dim=-1)
         crying_probabilities = probabilities[:, 1]
         avg_crying_probability = crying_probabilities.mean()*100
         if avg_crying_probability < 15:
@@ -54,11 +55,11 @@ def predict_stream(audio_path_stream):
             return "No está llorando"
 
 def decibelios(audio_path_stream):
-    waveform, sample_rate = torchaudio.load(audio_path_stream)
+    waveform, _ = torchaudio.load(audio_path_stream)
     rms = torch.sqrt(torch.mean(torch.square(waveform)))
     db_level = 20 * torch.log10(rms + 1e-6).item()
-    min_db = -80 
-    max_db = 0   
+    min_db = -80
+    max_db = 0
     scaled_db_level = (db_level - min_db) / (max_db - min_db)
     normalized_db_level = scaled_db_level * 100
     return normalized_db_level
@@ -66,15 +67,15 @@ def decibelios(audio_path_stream):
 def mostrar_decibelios(audio_path_stream, visual_threshold):
     db_level = decibelios(audio_path_stream)
     if db_level > visual_threshold:
-        return f"Prediciendo... Decibelios: {db_level:.2f}"
-    elif db_level < visual_threshold:
-        return f"Esperando... Decibelios: {db_level:.2f}"
+        status = "Prediciendo..."
+    else:
+        status = "Esperando..."
+    return f"<h3 style='text-align: center; font-size: 1.5em;'>{status} Decibelios: {db_level:.2f}</h3>"
 
 def predict_stream_decib(audio_path_stream, visual_threshold):
     db_level = decibelios(audio_path_stream)
     if db_level > visual_threshold:
-        llorando = predict_stream(audio_path_stream)
-        return f"{llorando}"
+        return display_prediction_stream(audio_path_stream)
     else:
         return ""
 
@@ -99,6 +100,16 @@ def chatbot_config(message, history: list[tuple[str, str]]):
 def cambiar_pestaña():
     return gr.update(visible=False), gr.update(visible=True)
 
+def display_prediction(audio, prediction_func):
+    prediction = prediction_func(audio)
+    return f"<h3 style='text-align: center; font-size: 1.5em;'>{prediction}</h3>"
+
+def display_prediction_wrapper(audio):
+    return display_prediction(audio, predict)
+
+def display_prediction_stream(audio):
+    return display_prediction(audio, predict_stream)
+
 my_theme = gr.themes.Soft(
     primary_hue="emerald",
     secondary_hue="green",
@@ -201,39 +212,41 @@ with gr.Blocks(theme=my_theme) as demo:
             label="Baby recorder",
             type="filepath",
             )
+        prediction_output = gr.Markdown()
         gr.Button("¿Por qué llora?").click(
-            predict,
+            display_prediction_wrapper,
             inputs=audio_input,
-            outputs=gr.Textbox(label="Tu bebé llora por:")
+            outputs=gr.Markdown()
             )
         gr.Button("Volver").click(cambiar_pestaña, outputs=[pag_predictor, chatbot])
     with gr.Column(visible=False) as pag_monitor:
         gr.Markdown("<h2>Monitor</h2>")
         gr.Markdown("<h4 style='text-align: center; font-size: 1.5em'>Detecta en tiempo real si tu bebé está llorando y por qué</h4>")
         audio_stream = gr.Audio(
-                format="wav",
-                label="Baby recorder",
-                type="filepath",
-                streaming=True
-            )
+            format="wav",
+            label="Baby recorder",
+            type="filepath",
+            streaming=True
+        )
         threshold_db = gr.Slider(
             minimum=0,
-            maximum=100,
+            maximum=120,
             step=1,
-            value=30,
+            value=20,
             label="Umbral de ruido para activar la predicción:"
-            )
+        )
+        volver = gr.Button("Volver")
         audio_stream.stream(
             mostrar_decibelios,
             inputs=[audio_stream, threshold_db],
-            outputs=gr.Textbox(value="Esperando...", label="Estado")
-            )
+            outputs=gr.HTML(value="<h3 style='text-align: center; font-size: 1.5em;'>Esperando...</h3>")
+        )
         audio_stream.stream(
             predict_stream_decib,
             inputs=[audio_stream, threshold_db],
-            outputs=gr.Textbox(value="", label="Tu bebé:")
+            outputs=gr.HTML()
         )
-        gr.Button("Volver").click(cambiar_pestaña, outputs=[pag_monitor, chatbot])
+        volver.click(cambiar_pestaña, outputs=[pag_monitor, chatbot])
     boton_inicial.click(cambiar_pestaña, outputs=[inicial, chatbot])
     boton_predictor.click(cambiar_pestaña, outputs=[chatbot, pag_predictor])
     boton_monitor.click(cambiar_pestaña, outputs=[chatbot, pag_monitor])