Update app.py

app.py

@@ -1,84 +1,76 @@
-import gradio as gr
-import numpy as np
 import tensorflow as tf
 import tensorflow_hub as hub
-import tensorflow_io as tfio
+import numpy as np
 import matplotlib.pyplot as plt
-import io
-from PIL import Image
-import pandas as pd
+import gradio as gr
+import os
+import scipy.io.wavfile as wavfile
 
-# Load YAMNet model
-yamnet_model_handle = 'https://tfhub.dev/google/yamnet/1'
+# Load YAMNet model from TensorFlow Hub
+yamnet_model_handle = "https://tfhub.dev/google/yamnet/1"
 yamnet_model = hub.load(yamnet_model_handle)
 
-# Load class names
-class_map_path = yamnet_model.class_map_path().numpy().decode('utf-8')
-class_names = list(pd.read_csv(class_map_path)['display_name'])
+# Load class names for YAMNet
+def load_class_map():
+    class_map_path = tf.keras.utils.get_file(
+        'yamnet_class_map.csv',
+        'https://raw.githubusercontent.com/tensorflow/models/master/research/audioset/yamnet/yamnet_class_map.csv'
+    )
+    with open(class_map_path, 'r') as f:
+        class_names = [line.strip().split(',')[2] for line in f.readlines()[1:]]
+    return class_names
+
+class_names = load_class_map()
 
-# Decode and resample audio
-def load_wav_16k_mono(audio_bytes):
-    audio_tensor, sample_rate = tf.audio.decode_wav(audio_bytes, desired_channels=1)
-    audio_tensor = tf.squeeze(audio_tensor, axis=-1)
-    audio_tensor = tfio.audio.resample(audio_tensor, rate_in=tf.cast(sample_rate, tf.int64), rate_out=16000)
-    return audio_tensor
+# Function to preprocess and classify audio
+def classify_audio(file_path):
+    try:
+        # Read audio file
+        sample_rate, audio_data = wavfile.read(file_path)
 
-# Plot waveform
-def plot_waveform(audio_tensor):
-    plt.figure(figsize=(8, 2))
-    plt.plot(audio_tensor.numpy())
-    plt.title("Waveform")
-    plt.tight_layout()
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    plt.close()
-    buf.seek(0)
-    return Image.open(buf)
+        # Ensure mono
+        if len(audio_data.shape) > 1:
+            audio_data = np.mean(audio_data, axis=1)
 
-# Plot log-mel spectrogram
-def plot_spectrogram(spectrogram):
-    plt.figure(figsize=(8, 3))
-    plt.imshow(spectrogram.numpy().T, aspect='auto', origin='lower', interpolation='nearest')
-    plt.title("Log-mel Spectrogram")
-    plt.xlabel("Frames")
-    plt.ylabel("Mel Bands")
-    plt.tight_layout()
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png')
-    plt.close()
-    buf.seek(0)
-    return Image.open(buf)
+        # Normalize audio
+        audio_data = audio_data / np.max(np.abs(audio_data))
 
-# Gradio interface logic
-def classify_sound(audio_file):
-    if isinstance(audio_file, str):
-        audio_bytes = tf.io.read_file(audio_file)
-    else:
-        audio_bytes = audio_file.read()
+        # Run inference
+        scores, embeddings, spectrogram = yamnet_model(audio_data)
+        scores_np = scores.numpy()
 
-    waveform = load_wav_16k_mono(audio_bytes)
-    scores, embeddings, spectrogram = yamnet_model(waveform)
+        # Get mean scores
+        mean_scores = np.mean(scores_np, axis=0)
+        top_5_indices = np.argsort(mean_scores)[::-1][:5]
+        top_class = class_names[top_5_indices[0]]
 
-    mean_scores = tf.reduce_mean(scores, axis=0)
-    top_class = tf.math.argmax(mean_scores)
-    inferred_class = class_names[top_class]
+        # Prepare waveform plot
+        fig, ax = plt.subplots()
+        ax.plot(audio_data)
+        ax.set_title("Waveform")
+        ax.set_xlabel("Sample Index")
+        ax.set_ylabel("Amplitude")
+        plt.tight_layout()
 
-    waveform_img = plot_waveform(waveform)
-    spectrogram_img = plot_spectrogram(spectrogram)
+        # Return predictions and plot
+        return top_class, {class_names[i]: float(mean_scores[i]) for i in top_5_indices}, fig
 
-    return inferred_class, waveform_img, spectrogram_img
+    except Exception as e:
+        return f"Error processing audio: {str(e)}", {}, None
 
-# Gradio app
-app = gr.Interface(
-    fn=classify_sound,
-    inputs=gr.Audio(type="file", label="Upload audio file"),
+# Build Gradio interface
+interface = gr.Interface(
+    fn=classify_audio,
+    inputs=gr.Audio(type="filepath", label="Upload .wav or .mp3 audio file"),
     outputs=[
-        gr.Text(label="Predicted Class"),
-        gr.Image(type="pil", label="Waveform"),
-        gr.Image(type="pil", label="Log-mel Spectrogram")
+        gr.Textbox(label="Top Prediction"),
+        gr.Label(label="Top 5 Classes with Scores"),
+        gr.Plot(label="Waveform")
     ],
-    title="YAMNet Audio Classifier",
-    description="Classify environmental and animal sounds using YAMNet. Visualize waveform and log-mel spectrogram."
+    title="Audtheia YAMNet Audio Classifier",
+    description="Upload an environmental or animal sound to classify using the YAMNet model. Returns label predictions and waveform."
 )
 
-app.launch()
+# Launch app
+if __name__ == "__main__":
+    interface.launch()