Update app.py

app.py

@@ -5,22 +5,20 @@ from app.preprocess import preprocess_audio
 from app.model import predict
 from collections import Counter, defaultdict
 import librosa
+import random
+from datasets import load_dataset
+
+# Load Hugging Face datasets
+audio_ds = load_dataset("AIOmarRehan/General_Audio_Dataset", split="train")
+image_ds = load_dataset("AIOmarRehan/Mel_Spectrogram_Images_for_Audio_Classification", split="train")
 
 
 # IMAGE HANDLING
 def safe_load_image(img):
-    """
-    Ensure the input is a valid PIL RGBA image.
-    Gradio sometimes gives numpy arrays → we convert safely.
-    """
     if img is None:
         return None
-
-    # If numpy array → convert to PIL
     if isinstance(img, np.ndarray):
         img = Image.fromarray(img)
-
-    # Convert to RGBA, to make sure the Alpha channel keep
     img = img.convert("RGBA")
     return img
 
@@ -34,12 +32,8 @@ def process_image_input(img):
 
 # PROCESS RAW AUDIO
 def process_audio_input(audio_path):
-
-    imgs = preprocess_audio(audio_path)  # returns list of PIL RGBA images
-
-    all_preds = []
-    all_confs = []
-    all_probs = []
+    imgs = preprocess_audio(audio_path)
+    all_preds, all_confs, all_probs = [], [], []
 
     for img in imgs:
         label, conf, probs = predict(img)
@@ -56,20 +50,28 @@ def process_audio_input(audio_path):
         final_label = candidates[0]
     else:
         conf_sums = defaultdict(float)
-        for i, label in enumerate(all_preds):
-            if label in candidates:
-                conf_sums[label] += all_confs[i]
+        for i, lbl in enumerate(all_preds):
+            if lbl in candidates:
+                conf_sums[lbl] += all_confs[i]
         final_label = max(conf_sums, key=conf_sums.get)
 
-    final_conf = float(
-        np.mean([all_confs[i] for i, lbl in enumerate(all_preds) if lbl == final_label])
-    )
+    final_conf = float(np.mean([all_confs[i] for i, lbl in enumerate(all_preds) if lbl == final_label]))
 
     return final_label, round(final_conf, 3), all_preds, [round(c, 3) for c in all_confs]
 
 
 # MAIN CLASSIFIER
-def classify(audio_path, image):
+def classify(audio_path, image, random_audio, random_image):
+    # Load random audio if selected
+    if random_audio:
+        rand_sample = random.choice(audio_ds)
+        audio_path = rand_sample["audio"]["path"]
+
+    # Load random image if selected
+    if random_image:
+        rand_sample = random.choice(image_ds)
+        img_bytes = rand_sample["image"]
+        image = Image.open(img_bytes).convert("RGBA")
 
     # If spectrogram image
     if image is not None:
@@ -78,7 +80,7 @@ def classify(audio_path, image):
             "Final Label": label,
             "Confidence": conf,
             "Details": probs
-        }
+        }, label  
 
     # If raw audio
     if audio_path is not None:
@@ -88,9 +90,9 @@ def classify(audio_path, image):
             "Confidence": conf,
             "All Chunk Labels": all_preds,
             "All Chunk Confidences": all_confs
-        }
+        }, label  
 
-    return "Please upload an audio file OR a spectrogram image."
+    return "Please upload an audio file OR a spectrogram image.", ""
 
 
 # GRADIO UI
@@ -98,14 +100,19 @@ interface = gr.Interface(
     fn=classify,
     inputs=[
         gr.Audio(type="filepath", label="Upload Audio (WAV/MP3)"),
-        gr.Image(type="pil", label="Upload Spectrogram Image (PNG RGBA Supported)")
+        gr.Image(type="pil", label="Upload Spectrogram Image (PNG RGBA Supported)"),
+        gr.Checkbox(label="Pick Random Audio from Dataset"),
+        gr.Checkbox(label="Pick Random Image from Dataset"),
+    ],
+    outputs=[
+        gr.JSON(label="Prediction Results"),
+        gr.Textbox(label="Final Label", interactive=False)
     ],
-    outputs=gr.JSON(label="Prediction Results"),
     title="General Audio Classifier (Audio + Spectrogram Support)",
     description=(
         "Upload a raw audio file OR a spectrogram image.\n"
-        "If audio → model preprocesses into mel-spectrogram chunks.\n"
-        "If image → model classifies the spectrogram directly.\n"
+        "You can also select random samples from the Hugging Face datasets.\n"
+        "The output shows a JSON with all details and a separate field for the final label."
     ),
 )