Updated with wandb settings

app.py

@@ -2,9 +2,14 @@ from fastai.vision.all import *
 import librosa
 import gradio as gr
 import numpy as np
+import wandb
+from fastai.callback.wandb import *
+
+os.environ["WANDB_API_KEY"] = WANDB_API_KEY
+wandb.init()
 
 learn = load_learner('audio_mnist_classifier_v1.pkl')
-categories = ('0', '1', '2', '3', '4', '5', '6', '7', '8', '9')
+categories = learn.dls.vocab
 
 def mel_spectrogram_tfm(file):
     y, sr = librosa.load(file)
@@ -15,8 +20,16 @@ def mel_spectrogram_tfm(file):
 
 def classify(audio):
     spec_db = mel_spectrogram_tfm(audio)
-    img = PILImage.create(spec_db).resize((225,225))
+    img = PILImage.create(spec_db)
     pred, idx, prob = learn.predict(img)
     return dict(zip(categories, map(float, prob)))
 
-gr.Interface(fn=classify, inputs=gr.Audio(source="microphone", type="filepath"), outputs="text", title="Identify digits between 0 and 9 from audio clips").launch()
+interface = gr.Interface(fn=classify, 
+             inputs=gr.Audio(source="microphone", type="filepath"), 
+             outputs=gr.outputs.Label(num_top_classes=10),
+             title='Audio MNIST Classification',
+             description='Identifying digits (from 0 to 9) from an audio clip')
+
+interface.launch()
+
+interface.integrate(wandb=wandb)