import gradio as gr
import torchaudio
from transformers import pipeline, AudioClassificationPipeline

pipe = pipeline("audio-classification", model="cogniveon/eeem069_heart_murmur_classification")

def predict(audio):
  sampling_rate, data = audio
  waveform = torch.tensor(data).float()

  # Resample the audio to 16 kHz (if necessary)
  if sampling_rate != 16000:
    resampler = torchaudio.transforms.Resample(sampling_rate, 16000)
    waveform = resampler(waveform)

  results = pipe(waveform.numpy())

  sorted_results = sorted(results, key=lambda x: x['score'], reverse=True)

  label_scores = {item['label']: item['score'] for item in sorted_results}

  return label_scores


gr.Interface(
    fn=predict,
    inputs="audio",
    outputs="label"
).launch()