Update app.py

app.py

@@ -26,9 +26,15 @@ def loading_sound_file(sound_file, sr=22050, duration=10):
     #     print ("fixing audio lenght :", file_name)
     #     y = lib.util.fix_length(X, input_length) 
     # extract normalized mfcc feature from data
+    
+    ## pad audio to same duration 
+    if round(dur) < duration:
+        X = lib.util.fix_length(X, input_length)
+    
     mfccs = np.mean(lib.feature.mfcc(y=X, sr=sr, n_mfcc=25).T,axis=0)
 
-    data = np.array(mfccs).reshape([-1,1])
+    ## Reshape to match the model's input shape
+    data = np.array(mfccs).reshape([1, -1, 1)
 
     return data
 
@@ -37,15 +43,26 @@ def loading_sound_file(sound_file, sr=22050, duration=10):
 def heart_signal_classification(data):
     X = loading_sound_file(data)
     pred = model.predict(X)
-    result = pred[0].argmax()
-    ## lets create our labels
+    ## Define the threshold
+    threshold = 0.6
+    max_prob = np.max(pred)
+    
+    ## Create labels
     labels = {
-        0: 'Artifact',
-        1: 'Murmur',
-        2: 'Normal' 
+        0: 'artifact',
+        1: 'unlabel',
+        2: 'extrastole',
+        3: 'extrahls',
+        4: 'normal',
+        5: 'murmur'
     }
-
-    label = labels[pred[0].argmax()]
+    
+    if max_prob < threshold:
+        label = 'unknown'
+    else:
+        result = pred[0].argmax()
+        label = labels[result]
+    
     return label
 ################### Gradio Web APP ################################ 
 title = "Heart Signal Classification App"