TODO FUNCIONANDO LOCAL

.gitignore

@@ -1,5 +1,5 @@
 __pycache__
 .venv
 .vscode
-A-POR-LOS-8000/distilhubert-finetuned-cry-detector
-A-POR-LOS-8000/distilhubert-finetuned-mixed-data
+distilhubert-finetuned-cry-detector
+distilhubert-finetuned-mixed-data

app.py

@@ -8,13 +8,23 @@ from interfaz import estilo, my_theme
 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")
-model_class, id2label_class = predict_params(model_path="A-POR-LOS-8000/distilhubert-finetuned-mixed-data", dataset_path="A-POR-LOS-8000/data/mixed_data", filter_white_noise=True)
-model_mon, id2label_mon = predict_params(model_path="A-POR-LOS-8000/distilhubert-finetuned-cry-detector", dataset_path="A-POR-LOS-8000/data/baby_cry_detection", filter_white_noise=False)
+model_class, id2label_class = predict_params(
+    model_path="distilhubert-finetuned-mixed-data",
+    dataset_path="data/mixed_data",
+    filter_white_noise=True,
+    undersample_normal=True
+    )
+model_mon, id2label_mon = predict_params(
+    model_path="distilhubert-finetuned-cry-detector",
+    dataset_path="data/baby_cry_detection",
+    filter_white_noise=False,
+    undersample_normal=False
+    )
 
-def call(audiopath, model, dataset_path, filter_white_noise):
+def call(audiopath, model, dataset_path, filter_white_noise, undersample_normal):
     model.to(device)
     model.eval()
-    audio_dataset = AudioDataset(dataset_path, {}, filter_white_noise,)
+    audio_dataset = AudioDataset(dataset_path, {}, filter_white_noise, undersample_normal)
     processed_audio = audio_dataset.preprocess_audio(audiopath)
     inputs = {"input_values": processed_audio.to(device).unsqueeze(0)}
     with torch.no_grad():
@@ -24,16 +34,16 @@ def call(audiopath, model, dataset_path, filter_white_noise):
 
 def predict(audio_path_pred):
     with torch.no_grad():
-        logits = call(audio_path_pred, model=model_class, dataset_path="A-POR-LOS-8000/data/mixed_data", filter_white_noise=True)
+        logits = call(audio_path_pred, model=model_class, dataset_path="data/mixed_data", filter_white_noise=True, undersample_normal=True)
         predicted_class_ids_class = torch.argmax(logits, dim=-1).item()
         label_class = id2label_class[predicted_class_ids_class]
-        label_mapping = {0: 'Hambre', 1: 'Problemas para respirar', 2: 'Dolor', 3: 'Cansancio/Incomodidad'}
+        label_mapping = {0: 'Dolor', 1: 'Cansancio/Incomodidad', 2: 'Hambre', 3: 'Problemas para respirar'}
         label_class = label_mapping.get(predicted_class_ids_class, label_class)
     return label_class
 
 def predict_stream(audio_path_stream):
     with torch.no_grad():
-        logits = call(audio_path_stream, model=model_mon, dataset_path="A-POR-LOS-8000/data/baby_cry_detection", filter_white_noise=False)
+        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)
         crying_probabilities = probabilities[:, 1]
         avg_crying_probability = crying_probabilities.mean()*100
@@ -45,7 +55,7 @@ def predict_stream(audio_path_stream):
 
 def decibelios(audio_path_stream):
     with torch.no_grad():
-        logits = call(audio_path_stream, model=model_mon, dataset_path="A-POR-LOS-8000/data/baby_cry_detection", filter_white_noise=False)
+        logits = call(audio_path_stream, model=model_mon, dataset_path="data/baby_cry_detection", filter_white_noise=False, undersample_normal=False)
         rms = torch.sqrt(torch.mean(torch.square(logits)))
         db_level = 20 * torch.log10(rms + 1e-6).item()
         return db_level