Update app.py

app.py

@@ -1,19 +1,22 @@
 import gradio as gr
 import numpy as np
 from PIL import Image
-import random
 import io
+import random
+import tempfile
 from collections import Counter, defaultdict
 from datasets import load_dataset
 from app.model import predict
 from app.preprocess import preprocess_audio
+import soundfile as sf
 
-# Load Hugging Face datasets directly
-audio_ds = load_dataset("AIOmarRehan/General_Audio_Dataset")
-image_ds = load_dataset("AIOmarRehan/Mel_Spectrogram_Images_for_Audio_Classification")
+# 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")
 
-# Helper function to safely load images
+# Helper functions
 def safe_load_image(img):
+    """Ensure input is PIL RGBA image"""
     if img is None:
         return None
     if isinstance(img, np.ndarray):
@@ -21,13 +24,11 @@ def safe_load_image(img):
     img = img.convert("RGBA")
     return img
 
-# Process spectrogram image
 def process_image_input(img):
     img = safe_load_image(img)
     label, confidence, probs = predict(img)
     return label, round(confidence, 3), probs
 
-# Process raw audio
 def process_audio_input(audio_path):
     imgs = preprocess_audio(audio_path)
     all_preds, all_confs, all_probs = [], [], []
@@ -55,28 +56,44 @@ def process_audio_input(audio_path):
     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
+# Main classifier function
 def classify(audio_path, image, random_audio=False, random_image=False):
-    # Pick random audio from HF dataset
+    # Random audio selection
     if random_audio and len(audio_ds) > 0:
-        sample = random.choice(audio_ds)
-        # If dataset stores audio as file path or array
-        if isinstance(sample["audio"], dict) and "path" in sample["audio"]:
-            audio_path = sample["audio"]["path"]
-        elif isinstance(sample["audio"], dict) and "array" in sample["audio"]:
-            # Save array temporarily
-            import soundfile as sf
-            audio_path = "/tmp/random_audio.wav"
-            sf.write(audio_path, sample["audio"]["array"], sample["audio"]["sampling_rate"])
+        try:
+            sample = random.choice(audio_ds)
+            # Dataset may store audio as path or array
+            audio_obj = sample["audio"]
+            if isinstance(audio_obj, dict) and "path" in audio_obj:
+                audio_path = audio_obj["path"]
+            elif isinstance(audio_obj, dict) and "array" in audio_obj:
+                # Save temporarily
+                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmpfile:
+                    audio_path = tmpfile.name
+                    sf.write(audio_path, audio_obj["array"], audio_obj["sampling_rate"])
+            else:
+                # fallback: datasets.Audio object
+                audio_array, sr = audio_obj["array"], audio_obj["sampling_rate"]
+                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmpfile:
+                    audio_path = tmpfile.name
+                    sf.write(audio_path, audio_array, sr)
+        except Exception as e:
+            print("Error loading random audio:", e)
+            audio_path = None
 
-    # Pick random image from HF dataset
+    # Random image selection
     if random_image and len(image_ds) > 0:
-        sample = random.choice(image_ds)
-        # Handle image bytes
-        img_bytes = sample["image"] if isinstance(sample["image"], bytes) else sample["image"].tobytes()
-        image = Image.open(io.BytesIO(img_bytes)).convert("RGBA")
+        try:
+            sample = random.choice(image_ds)
+            img_obj = sample["image"]
+            if not isinstance(img_obj, Image.Image):
+                img_obj = Image.fromarray(img_obj)  # convert ndarray to PIL
+            image = img_obj.convert("RGBA")
+        except Exception as e:
+            print("Error loading random image:", e)
+            image = None
 
-    # If spectrogram image
+    # Process spectrogram image
     if image is not None:
         label, conf, probs = process_image_input(image)
         return {
@@ -85,7 +102,7 @@ def classify(audio_path, image, random_audio=False, random_image=False):
             "Details": probs
         }, label
 
-    # If raw audio
+    # Process raw audio
     if audio_path is not None:
         label, conf, all_preds, all_confs = process_audio_input(audio_path)
         return {