Update app.py

app.py

@@ -5,7 +5,6 @@ import librosa
 import soundfile as sf
 from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
 import matplotlib.pyplot as plt
-from matplotlib.colors import Normalize
 import tempfile
 import os
 
@@ -14,63 +13,62 @@ SAMPLING_RATE = 16000
 MODEL_NAME = "MIT/ast-finetuned-audioset-10-10-0.4593"
 DEFAULT_THRESHOLD = 0.7
 
-# Load model and feature extractor
-feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
-model = AutoModelForAudioClassification.from_pretrained(MODEL_NAME)
+# Load model components
+try:
+    feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
+    model = AutoModelForAudioClassification.from_pretrained(MODEL_NAME)
+except Exception as e:
+    print(f"Error loading model: {str(e)}")
 
 # Equipment knowledge base
 EQUIPMENT_RECOMMENDATIONS = {
     "bearing": {
-        "high_frequency": "Recommend bearing replacement. High-frequency noise indicates wear or lubrication issues.",
-        "low_frequency": "Check for improper installation or contamination in bearings.",
-        "irregular": "Possible bearing cage damage. Schedule vibration analysis."
+        "high_frequency": "• Replace bearings immediately\n• Check lubrication system\n• Monitor vibration levels",
+        "low_frequency": "• Inspect bearing installation\n• Check for contamination\n• Verify lubrication",
+        "irregular": "• Perform vibration analysis\n• Schedule bearing replacement\n• Check alignment"
     },
     "pump": {
-        "cavitation": "Pump cavitation detected. Check suction conditions and NPSH.",
-        "impeller": "Impeller damage likely. Inspect and balance if needed.",
-        "misalignment": "Misalignment detected. Perform laser shaft alignment."
+        "cavitation": "• Check NPSH available\n• Inspect suction strainer\n• Adjust operating speed",
+        "impeller": "• Inspect impeller for damage\n• Perform dynamic balancing\n• Check wear rings",
+        "misalignment": "• Perform laser alignment\n• Check coupling condition\n• Verify baseplate level"
     },
     "motor": {
-        "electrical": "Electrical fault suspected. Check windings and connections.",
-        "mechanical": "Mechanical imbalance detected. Perform dynamic balancing.",
-        "bearing": "Motor bearing wear detected. Schedule replacement."
-    },
-    "compressor": {
-        "valve": "Compressor valve leakage suspected. Perform valve test.",
-        "pulsation": "Pulsation issues detected. Check dampeners and piping.",
-        "surge": "Compressor surge condition. Review control settings."
+        "electrical": "• Megger test windings\n• Check connections\n• Inspect starter contacts",
+        "mechanical": "• Perform dynamic balancing\n• Check alignment\n• Inspect cooling fins",
+        "bearing": "• Replace motor bearings\n• Check lubrication\n• Monitor temperature"
     }
 }
 
 def analyze_frequency_patterns(audio, sr):
     """Analyze frequency patterns to identify potential issues"""
     patterns = []
+    features = {}
     
     # Spectral analysis
     spectral_centroid = librosa.feature.spectral_centroid(y=audio, sr=sr)[0]
     spectral_rolloff = librosa.feature.spectral_rolloff(y=audio, sr=sr)[0]
     
-    mean_centroid = np.mean(spectral_centroid)
-    mean_rolloff = np.mean(spectral_rolloff)
+    features['centroid_mean'] = np.mean(spectral_centroid)
+    features['rolloff_mean'] = np.mean(spectral_rolloff)
     
-    if mean_centroid > 3000:  # High frequency components
+    if features['centroid_mean'] > 3000:
         patterns.append("high_frequency")
-    elif mean_centroid < 1000:  # Low frequency components
+    elif features['centroid_mean'] < 1000:
         patterns.append("low_frequency")
         
-    if mean_rolloff > 8000:  # Rich in harmonics
+    if features['rolloff_mean'] > 8000:
         patterns.append("harmonic_rich")
-        
-    return patterns
+    
+    return patterns, features
 
 def generate_recommendation(prediction, confidence, audio, sr):
     """Generate maintenance recommendations based on analysis"""
     if prediction == "Normal":
-        return "No immediate action required. Equipment operating within normal parameters."
+        return "✅ No immediate action required. Equipment operating within normal parameters."
     
-    patterns = analyze_frequency_patterns(audio, sr)
+    patterns, features = analyze_frequency_patterns(audio, sr)
     
-    # Simple equipment type classifier based on frequency profile
+    # Equipment classification
     spectral_flatness = librosa.feature.spectral_flatness(y=audio)[0]
     mean_flatness = np.mean(spectral_flatness)
     
@@ -79,78 +77,77 @@ def generate_recommendation(prediction, confidence, audio, sr):
     elif 0.2 <= mean_flatness < 0.6:
         equipment_type = "pump"
     else:
-        equipment_type = "motor" if np.mean(audio) < 0.1 else "compressor"
+        equipment_type = "motor"
     
-    # Generate specific recommendations
-    recommendations = ["🔧 Maintenance Recommendations:"]
-    recommendations.append(f"Detected issues in {equipment_type} with {confidence:.1%} confidence")
+    # Generate recommendations
+    recommendations = [
+        "🔧 MAINTENANCE RECOMMENDATIONS",
+        f"Equipment Type: {equipment_type.upper()}",
+        f"Confidence: {confidence:.1%}",
+        ""
+    ]
     
     for pattern in patterns:
         if pattern in EQUIPMENT_RECOMMENDATIONS.get(equipment_type, {}):
-            recommendations.append(f"→ {EQUIPMENT_RECOMMENDATIONS[equipment_type][pattern]}")
+            recommendations.append(EQUIPMENT_RECOMMENDATIONS[equipment_type][pattern])
     
-    # General recommendations
     if prediction == "Anomaly":
-        recommendations.append("\n🛠️ Suggested Actions:")
-        recommendations.append("1. Isolate equipment if possible")
-        recommendations.append("2. Perform visual inspection")
-        recommendations.append("3. Schedule detailed diagnostics")
-        recommendations.append(f"4. Review last maintenance records ({equipment_type})")
+        recommendations.extend([
+            "",
+            "🛠️ GENERAL ACTIONS:",
+            "1. Isolate equipment if possible",
+            "2. Perform visual inspection",
+            "3. Schedule detailed diagnostics",
+        ])
         
         if confidence > 0.8:
-            recommendations.append("\n🚨 Urgent: High confidence abnormality detected. Recommend immediate inspection!")
-    
+            recommendations.append("\n🚨 URGENT: High-confidence abnormality detected!")
+
     return "\n".join(recommendations)
 
+def process_audio(file_path):
+    """Handle audio file processing"""
+    try:
+        audio, sr = librosa.load(file_path, sr=SAMPLING_RATE, mono=True)
+        return audio, sr
+    except Exception as e:
+        raise RuntimeError(f"Audio processing error: {str(e)}")
+
 def analyze_audio(audio_input, threshold=DEFAULT_THRESHOLD):
-    """Process audio and detect anomalies"""
+    """Main analysis function"""
     try:
         # Handle file upload
         if isinstance(audio_input, str):
-            audio, sr = sf.read(audio_input)
-        else:  # Gradio file object
+            audio, sr = process_audio(audio_input)
+        else:  # Handle file object
             with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp:
                 tmp.write(audio_input.read())
                 tmp_path = tmp.name
-            audio, sr = sf.read(tmp_path)
+            audio, sr = process_audio(tmp_path)
             os.unlink(tmp_path)
         
-        # Convert to mono and resample if needed
-        if len(audio.shape) > 1:
-            audio = np.mean(audio, axis=1)
-        if sr != SAMPLING_RATE:
-            audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLING_RATE)
-        
-        # Feature extraction and prediction
+        # Model prediction
         inputs = feature_extractor(audio, sampling_rate=SAMPLING_RATE, return_tensors="pt")
         with torch.no_grad():
             outputs = model(**inputs)
             probs = torch.softmax(outputs.logits, dim=-1)
         
-        # Get results
         predicted_class = "Normal" if probs[0][0] > threshold else "Anomaly"
         confidence = probs[0][0].item() if predicted_class == "Normal" else 1 - probs[0][0].item()
         
-        # Generate spectrogram
-        spectrogram = librosa.feature.melspectrogram(y=audio, sr=SAMPLING_RATE, n_mels=64, fmax=8000)
-        db_spec = librosa.power_to_db(spectrogram, ref=np.max)
-        
-        fig, ax = plt.subplots(figsize=(10, 4))
-        librosa.display.specshow(db_spec, x_axis='time', y_axis='mel', sr=SAMPLING_RATE, fmax=8000, ax=ax)
+        # Generate visualization
+        plt.figure(figsize=(10, 4))
+        S = librosa.feature.melspectrogram(y=audio, sr=SAMPLING_RATE, n_mels=64)
+        S_db = librosa.power_to_db(S, ref=np.max)
+        librosa.display.specshow(S_db, x_axis='time', y_axis='mel', sr=SAMPLING_RATE, fmax=8000)
         plt.colorbar(format='%+2.0f dB')
-        plt.title('Mel Spectrogram with Anomaly Detection')
-        
-        # Mark anomalies on plot
-        if predicted_class == "Anomaly":
-            plt.text(0.5, 0.9, 'ANOMALY DETECTED', color='red', 
-                    ha='center', va='center', transform=ax.transAxes,
-                    fontsize=14, bbox=dict(facecolor='white', alpha=0.8))
+        plt.title('Mel Spectrogram')
         
         spec_path = os.path.join(tempfile.gettempdir(), 'spec.png')
         plt.savefig(spec_path, bbox_inches='tight')
         plt.close()
         
-        # Generate detailed recommendations
+        # Generate recommendations
         recommendations = generate_recommendation(predicted_class, confidence, audio, SAMPLING_RATE)
         
         return (
@@ -159,36 +156,35 @@ def analyze_audio(audio_input, threshold=DEFAULT_THRESHOLD):
             spec_path,
             recommendations
         )
-        
+    
     except Exception as e:
         return f"Error: {str(e)}", "", None, ""
 
 # Gradio Interface
-with gr.Blocks(title="Industrial Diagnostic Assistant 👨��🔧", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Industrial Audio Analyzer", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🏭 Industrial Equipment Diagnostic Assistant
-    ## Acoustic Anomaly Detection & Maintenance Recommendation System
+    # 🏭 Industrial Equipment Sound Analyzer
+    ### Acoustic Anomaly Detection & Maintenance Recommendation System
     """)
     
     with gr.Row():
         with gr.Column():
             audio_input = gr.Audio(
-                label="Upload Equipment Recording (.wav)",
-                type="filepath",
-                source="upload"
+                label="Upload Equipment Audio (.wav)",
+                type="filepath"
             )
             threshold = gr.Slider(
                 minimum=0.5, maximum=0.95, step=0.05, value=DEFAULT_THRESHOLD,
-                label="Detection Sensitivity", interactive=True
+                label="Detection Sensitivity"
             )
             analyze_btn = gr.Button("🔍 Analyze & Diagnose", variant="primary")
             
         with gr.Column():
             result_label = gr.Label(label="Diagnosis Result")
             confidence = gr.Textbox(label="Confidence Score")
-            spectrogram = gr.Image(label="Acoustic Analysis")
+            spectrogram = gr.Image(label="Spectrogram Analysis")
             recommendations = gr.Textbox(
-                label="Maintenance Recommendations", 
+                label="Maintenance Recommendations",
                 lines=10,
                 interactive=False
             )
@@ -200,15 +196,14 @@ with gr.Blocks(title="Industrial Diagnostic Assistant 👨‍🔧", theme=gr.the
     )
     
     gr.Markdown("""
-    ### System Capabilities:
-    - Automatic anomaly detection in industrial equipment sounds
-    - Frequency pattern analysis to identify failure modes
-    - Equipment-specific maintenance recommendations
-    - Confidence-based urgency classification
-    
-    **Tip:** For best results, use 5-10 second recordings of steady operation
+    **Instructions:**
+    - Upload 5-10 second .wav recordings
+    - Results include:
+      ✓ Anomaly detection
+      ✓ Equipment classification
+      ✓ Maintenance recommendations
+      ✓ Spectrogram visualization
     """)
 
 if __name__ == "__main__":
     demo.launch()
-