Update app.py

app.py

@@ -8,15 +8,12 @@ import matplotlib.pyplot as plt
 import tensorflow as tf
 from tensorflow.keras import models
 
-# --- 1. Configuration & Global Variables ---
-# Create a temporary directory for spectrograms if it doesn't exist
+# ---------------- 1. Configuration ---------------- #
 TEMP_DIR = "temp_gradio_specs"
 os.makedirs(TEMP_DIR, exist_ok=True)
-
-# Define image size for the model
 IMG_SIZE = (224, 224)
 
-# --- 2. Load Models and Define Classes (Done once on startup) ---
+# ---------------- 2. Load Models ------------------ #
 print("🚀 Loading machine learning models...")
 try:
     stage1_model = models.load_model("saved_models/stage1_model.h5")
@@ -25,41 +22,55 @@ try:
     print("✅ Models loaded successfully.")
 except Exception as e:
     print(f"❌ Error loading models: {e}")
-    # Exit if models can't be loaded
-    exit()
+    # Do not exit—allows app to show error gracefully
+    stage1_model = abnormal_model = normal_model = None
 
-# Define class lists exactly as they were during training
+# Default class lists – replace with actual labels if available
 stage1_classes = ["00 - Abnormal", "01 - Normal"]
-abnormal_classes = sorted(os.listdir("MelSpectrograms/00 - Abnormal"))
-normal_classes = sorted(os.listdir("MelSpectrograms/01 - Normal"))
+
+abnormal_classes = (
+    sorted(os.listdir("MelSpectrograms/00 - Abnormal"))
+    if os.path.exists("MelSpectrograms/00 - Abnormal")
+    else ["Bearing noise", "Dehydration mode noise"]
+)
+
+normal_classes = (
+    sorted(os.listdir("MelSpectrograms/01 - Normal"))
+    if os.path.exists("MelSpectrograms/01 - Normal")
+    else ["Wash mode", "Spin mode"]
+)
 
 print(f"Stage 1 Classes: {stage1_classes}")
 print(f"Abnormal Sub-classes: {abnormal_classes}")
 print(f"Normal Sub-classes: {normal_classes}")
 
 
-# --- 3. Helper Functions and Classes ---
-
-def save_mel_spectrogram(file_path, save_dir, sr=22050, n_mels=128, hop_length=512, n_fft=2048):
+# ---------------- 3. Helper Functions -------------- #
+def save_mel_spectrogram(file_path, save_dir, sr=22050,
+                         n_mels=128, hop_length=512, n_fft=2048):
     """Generates and saves a Mel Spectrogram from an audio file."""
     try:
         y, sr = librosa.load(file_path, sr=sr, mono=True)
-        S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length)
+        S = librosa.feature.melspectrogram(
+            y=y, sr=sr, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length
+        )
         S_db = librosa.power_to_db(S, ref=np.max)
-        
+
         filename = os.path.basename(file_path).replace(".wav", ".png")
         save_path = os.path.join(save_dir, filename)
 
         plt.figure(figsize=(4, 4))
-        librosa.display.specshow(S_db, sr=sr, hop_length=hop_length, x_axis='time', y_axis='mel', cmap='magma')
+        librosa.display.specshow(S_db, sr=sr, hop_length=hop_length,
+                                 x_axis="time", y_axis="mel", cmap="magma")
         plt.axis("off")
         plt.savefig(save_path, bbox_inches="tight", pad_inches=0)
         plt.close()
         return save_path
     except Exception as e:
-        print(f"Error creating spectrogram: {e}")
+        print(f"❌ Error creating spectrogram: {e}")
         return None
 
+
 class HierarchicalClassifier:
     """A wrapper class for the two-stage prediction logic."""
     def __init__(self, stage1_model, abnormal_model, normal_model,
@@ -75,10 +86,18 @@ class HierarchicalClassifier:
     def _preprocess_image(self, image_path):
         img = tf.keras.utils.load_img(image_path, target_size=self.img_size)
         img_array = tf.keras.utils.img_to_array(img) / 255.0
-        img_array = tf.expand_dims(img_array, 0)
-        return img_array
+        return tf.expand_dims(img_array, 0)
 
     def predict(self, image_path):
+        if not all([self.stage1_model, self.abnormal_model, self.normal_model]):
+            return {
+                "final_prediction": "❌ Models not loaded. Please upload models to /saved_models/",
+                "stage1_class": "N/A",
+                "stage1_confidence": 0,
+                "stage2_class": "N/A",
+                "stage2_confidence": 0
+            }
+
         img_array = self._preprocess_image(image_path)
         stage1_pred = self.stage1_model.predict(img_array, verbose=0)
         stage1_idx = np.argmax(stage1_pred)
@@ -98,94 +117,59 @@ class HierarchicalClassifier:
             "stage1_confidence": float(np.max(stage1_pred)),
             "stage2_class": sub_class,
             "stage2_confidence": float(np.max(sub_pred)),
-            "final_prediction": f"{main_class.split(' - ')[1]} → {sub_class.split(' - ')[1]}"
+            "final_prediction": f"{main_class.split(' - ')[1]} → {sub_class}"
         }
 
-# Instantiate the classifier with loaded models and classes
+
 classifier = HierarchicalClassifier(
     stage1_model, abnormal_model, normal_model,
     stage1_classes, abnormal_classes, normal_classes
 )
 
-# --- 4. The Main Prediction Function for Gradio ---
+
+# ---------------- 4. Prediction Function ----------- #
 def predict_washing_machine_sound(audio_filepath):
-    """
-    This is the core function that Gradio will call.
-    It takes an audio file path, processes it, and returns the formatted result.
-    """
     if audio_filepath is None:
         return "Please upload an audio file first.", None
 
     print(f"Processing file: {audio_filepath}")
-    
-    # The spectrogram path needs to be cleaned up after prediction
-    spec_path = None
-    try:
-        # Generate a spectrogram from the input audio file
-        spec_path = save_mel_spectrogram(audio_filepath, TEMP_DIR)
-        if not spec_path:
-            return "Error: Could not generate spectrogram from the audio file.", None
-        
-        # Get prediction from the classifier
-        result = classifier.predict(spec_path)
-        
-        # Format the output for better readability
-        output_text = (
-            f"🎯 Final Prediction: {result['final_prediction']}\n\n"
-            f"Confidence Scores:\n"
-            f"--------------------\n"
-            f"Stage 1 ({result['stage1_class']}): {result['stage1_confidence']:.4f}\n"
-            f"Stage 2 ({result['stage2_class']}): {result['stage2_confidence']:.4f}"
-        )
-        
-        # Return the formatted text and the path to the spectrogram image to display it
-        return output_text, spec_path
+    spec_path = save_mel_spectrogram(audio_filepath, TEMP_DIR)
+    if not spec_path:
+        return "❌ Could not generate spectrogram from the audio file.", None
+
+    result = classifier.predict(spec_path)
+
+    output_text = (
+        f"🎯 Final Prediction: {result['final_prediction']}\n\n"
+        f"Confidence Scores:\n"
+        f"--------------------\n"
+        f"Stage 1 ({result['stage1_class']}): {result['stage1_confidence']:.4f}\n"
+        f"Stage 2 ({result['stage2_class']}): {result['stage2_confidence']:.4f}"
+    )
 
-    except Exception as e:
-        print(f"An error occurred during prediction: {e}")
-        return f"An error occurred: {str(e)}", None
-        
-    finally:
-        # Clean up the generated spectrogram image file after it's been used
-        # Gradio handles the temp audio file, but we must handle the temp spectrogram
-        if spec_path and os.path.exists(spec_path):
-            # Note: Gradio might need the file to display it, so cleaning up here
-            # might be too early if the image component relies on the path.
-            # For simplicity, we can let them accumulate in the temp folder or
-            # implement more complex cleanup later. Let's comment out the immediate delete.
-            # os.remove(spec_path)
-            pass
-
-# --- 5. Build and Launch the Gradio Interface ---
-if __name__ == "__main__":
-    # Define some example audio files from your dataset
-    example_files = [
-        "Washing machine/00 - Abnormal/00-2 - Dehydration mode noise/04.wav",
-        "Washing machine/01 - Normal/01-1 - Washing mode/01.wav",
-        "Washing machine/00 - Abnormal/00-1 - Bearing noise/02.wav"
-    ]
+    return output_text, spec_path
 
+
+# ---------------- 5. Gradio Interface -------------- #
+if __name__ == "__main__":
     demo = gr.Interface(
         fn=predict_washing_machine_sound,
-        inputs=gr.Audio(type="filepath", label="Upload Washing Machine Audio (.wav)"),
+        inputs=gr.Audio(type="filepath", label="Upload Washing-Machine Audio (.wav)"),
         outputs=[
             gr.Textbox(label="Prediction Result"),
-            gr.Image(label="Generated Mel Spectrogram")
+            gr.Image(label="Generated Mel-Spectrogram")
         ],
-        title="Washing Machine Sound Classifier",
-        description="Upload a WAV audio file of a washing machine to classify its operation status. The model performs a two-stage classification: first identifying 'Normal' vs 'Abnormal' sound, then determining the specific sub-type.",
-        # examples=example_files,
-        allow_flagging="never"
+        title="Washing-Machine Sound Classifier",
+        description="Upload a WAV file of washing-machine audio to classify its operation status.",
+        allow_flagging="never",
+        # examples=[]  # ← removed local file examples
     )
 
-    # Launch the web UI
     demo.launch()
 
-    # Clean up the entire temp directory on exit
-    # This is a simple way to manage temp files
+    # Cleanup temp dir after app stops
     try:
-        print("\nCleaning up temporary files...")
         shutil.rmtree(TEMP_DIR)
-        print("✅ Cleanup complete.")
+        print("✅ Cleaned up temporary files.")
     except Exception as e:
-        print(f"Could not clean up temp files: {e}")
+        print(f"⚠️ Cleanup warning: {e}")