import gradio as gr
import torch
import torchaudio
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification

# הגדרת המודל
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_file):
    # טען את קובץ הסאונד
    waveform, sample_rate = torchaudio.load(audio_file)

    # המרה ל־16kHz כפי שהמודל דורש
    if sample_rate != 16000:
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
        waveform = resampler(waveform)
        sample_rate = 16000

    # אם האודיו בסטראו, נאחד לערוץ אחד
    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="המערכת מאזינה לקובץ קול של תינוק ומחזירה את סוג הבכי לפי המודל"
).launch()