zainulabedin949/Audio-Spectrogram-Transformer

1import gradio as gr2import numpy as np3import torch4import librosa5import soundfile as sf6from transformers import AutoFeatureExtractor, AutoModelForAudioClassification7import matplotlib.pyplot as plt8import tempfile9import os10 11# Constants12SAMPLING_RATE = 1600013MODEL_NAME = "MIT/ast-finetuned-audioset-10-10-0.4593"14DEFAULT_THRESHOLD = 0.715 16# Load model components17try:18    feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)19    model = AutoModelForAudioClassification.from_pretrained(MODEL_NAME)20except Exception as e:21    print(f"Error loading model: {str(e)}")22 23# Equipment knowledge base24EQUIPMENT_RECOMMENDATIONS = {25    "bearing": {26        "high_frequency": "• Replace bearings immediately\n• Check lubrication system\n• Monitor vibration levels",27        "low_frequency": "• Inspect bearing installation\n• Check for contamination\n• Verify lubrication",28        "irregular": "• Perform vibration analysis\n• Schedule bearing replacement\n• Check alignment"29    },30    "pump": {31        "cavitation": "• Check NPSH available\n• Inspect suction strainer\n• Adjust operating speed",32        "impeller": "• Inspect impeller for damage\n• Perform dynamic balancing\n• Check wear rings",33        "misalignment": "• Perform laser alignment\n• Check coupling condition\n• Verify baseplate level"34    },35    "motor": {36        "electrical": "• Megger test windings\n• Check connections\n• Inspect starter contacts",37        "mechanical": "• Perform dynamic balancing\n• Check alignment\n• Inspect cooling fins",38        "bearing": "• Replace motor bearings\n• Check lubrication\n• Monitor temperature"39    }40}41 42def analyze_frequency_patterns(audio, sr):43    """Analyze frequency patterns to identify potential issues"""44    patterns = []45    features = {}46    47    # Spectral analysis48    spectral_centroid = librosa.feature.spectral_centroid(y=audio, sr=sr)[0]49    spectral_rolloff = librosa.feature.spectral_rolloff(y=audio, sr=sr)[0]50    51    features['centroid_mean'] = np.mean(spectral_centroid)52    features['rolloff_mean'] = np.mean(spectral_rolloff)53    54    if features['centroid_mean'] > 3000:55        patterns.append("high_frequency")56    elif features['centroid_mean'] < 1000:57        patterns.append("low_frequency")58        59    if features['rolloff_mean'] > 8000:60        patterns.append("harmonic_rich")61    62    return patterns, features63 64def generate_recommendation(prediction, confidence, audio, sr):65    """Generate maintenance recommendations based on analysis"""66    if prediction == "Normal":67        return "✅ No immediate action required. Equipment operating within normal parameters."68    69    patterns, features = analyze_frequency_patterns(audio, sr)70    71    # Equipment classification72    spectral_flatness = librosa.feature.spectral_flatness(y=audio)[0]73    mean_flatness = np.mean(spectral_flatness)74    75    if mean_flatness < 0.2:76        equipment_type = "bearing"77    elif 0.2 <= mean_flatness < 0.6:78        equipment_type = "pump"79    else:80        equipment_type = "motor"81    82    # Generate recommendations83    recommendations = [84        "🔧 MAINTENANCE RECOMMENDATIONS",85        f"Equipment Type: {equipment_type.upper()}",86        f"Confidence: {confidence:.1%}",87        ""88    ]89    90    for pattern in patterns:91        if pattern in EQUIPMENT_RECOMMENDATIONS.get(equipment_type, {}):92            recommendations.append(EQUIPMENT_RECOMMENDATIONS[equipment_type][pattern])93    94    if prediction == "Anomaly":95        recommendations.extend([96            "",97            "🛠️ GENERAL ACTIONS:",98            "1. Isolate equipment if possible",99            "2. Perform visual inspection",100            "3. Schedule detailed diagnostics",101        ])102        103        if confidence > 0.8:104            recommendations.append("\n🚨 URGENT: High-confidence abnormality detected!")105 106    return "\n".join(recommendations)107 108def process_audio(file_path):109    """Handle audio file processing"""110    try:111        audio, sr = librosa.load(file_path, sr=SAMPLING_RATE, mono=True)112        return audio, sr113    except Exception as e:114        raise RuntimeError(f"Audio processing error: {str(e)}")115 116def analyze_audio(audio_input, threshold=DEFAULT_THRESHOLD):117    """Main analysis function"""118    try:119        # Handle file upload120        if isinstance(audio_input, str):121            audio, sr = process_audio(audio_input)122        else:  # Handle file object123            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp:124                tmp.write(audio_input.read())125                tmp_path = tmp.name126            audio, sr = process_audio(tmp_path)127            os.unlink(tmp_path)128        129        # Model prediction130        inputs = feature_extractor(audio, sampling_rate=SAMPLING_RATE, return_tensors="pt")131        with torch.no_grad():132            outputs = model(**inputs)133            probs = torch.softmax(outputs.logits, dim=-1)134        135        predicted_class = "Normal" if probs[0][0] > threshold else "Anomaly"136        confidence = probs[0][0].item() if predicted_class == "Normal" else 1 - probs[0][0].item()137        138        # Generate visualization139        plt.figure(figsize=(10, 4))140        S = librosa.feature.melspectrogram(y=audio, sr=SAMPLING_RATE, n_mels=64)141        S_db = librosa.power_to_db(S, ref=np.max)142        librosa.display.specshow(S_db, x_axis='time', y_axis='mel', sr=SAMPLING_RATE, fmax=8000)143        plt.colorbar(format='%+2.0f dB')144        plt.title('Mel Spectrogram')145        146        spec_path = os.path.join(tempfile.gettempdir(), 'spec.png')147        plt.savefig(spec_path, bbox_inches='tight')148        plt.close()149        150        # Generate recommendations151        recommendations = generate_recommendation(predicted_class, confidence, audio, SAMPLING_RATE)152        153        return (154            predicted_class,155            f"{confidence:.1%}",156            spec_path,157            recommendations158        )159    160    except Exception as e:161        return f"Error: {str(e)}", "", None, ""162 163# Gradio Interface164with gr.Blocks(title="Industrial Audio Analyzer", theme=gr.themes.Soft()) as demo:165    gr.Markdown("""166    # 🏭 Industrial Equipment Sound Analyzer167    ### Acoustic Anomaly Detection & Maintenance Recommendation System168    """)169    170    with gr.Row():171        with gr.Column():172            audio_input = gr.Audio(173                label="Upload Equipment Audio (.wav)",174                type="filepath"175            )176            threshold = gr.Slider(177                minimum=0.5, maximum=0.95, step=0.05, value=DEFAULT_THRESHOLD,178                label="Detection Sensitivity"179            )180            analyze_btn = gr.Button("🔍 Analyze & Diagnose", variant="primary")181            182        with gr.Column():183            result_label = gr.Label(label="Diagnosis Result")184            confidence = gr.Textbox(label="Confidence Score")185            spectrogram = gr.Image(label="Spectrogram Analysis")186            recommendations = gr.Textbox(187                label="Maintenance Recommendations",188                lines=10,189                interactive=False190            )191    192    analyze_btn.click(193        fn=analyze_audio,194        inputs=[audio_input, threshold],195        outputs=[result_label, confidence, spectrogram, recommendations]196    )197    198    gr.Markdown("""199    **Instructions:**200    - Upload 5-10 second .wav recordings201    - Results include:202      ✓ Anomaly detection203      ✓ Equipment classification204      ✓ Maintenance recommendations205      ✓ Spectrogram visualization206    """)207 208if __name__ == "__main__":209    demo.launch()210