Update app.py

app.py

@@ -1,16 +1,32 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+import torchaudio
+from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
 
-classifier = pipeline("audio-classification", model="Wiam/baby-cry-classification-finetuned-babycry-v4")
+model_id = "Wiam/baby-cry-classification-finetuned-babycry-v4"
+model = AutoModelForAudioClassification.from_pretrained(model_id)
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
 
-def classify_baby_cry(audio):
-    results = classifier(audio)
-    return {res["label"]: res["score"] for res in results}
+def classify_baby_cry(audio_file):
+    waveform, sample_rate = torchaudio.load(audio_file)
+    
+    # אם האודיו ביותר מערוץ אחד (Stereo) נחתוך לערוץ 1
+    if waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+    
+    inputs = feature_extractor(waveform.squeeze().numpy(), sampling_rate=sample_rate, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = torch.nn.functional.softmax(outputs.logits[0], dim=0)
+
+    labels = model.config.id2label
+    results = {labels[i]: float(probs[i]) for i in range(len(labels))}
+    return results
 
 gr.Interface(
     fn=classify_baby_cry,
     inputs=gr.Audio(type="filepath"),
     outputs=gr.Label(num_top_classes=3),
     title="זיהוי בכי תינוקות",
-    description="העלה קובץ קול של בכי תינוקות ונדע לזהות מה הוא רוצה"
+    description="המערכת מאזינה לקובץ קול ומחזירה מה נראה שהתינוק מנסה לבטא"
 ).launch()