CNN_space / app.py
shiko217's picture
refactor
c942b33 unverified
Raw
History Blame Contribute Delete
2.32 kB
import os
import numpy as np
import joblib
import gradio as gr
import warnings
from ml_features import preprocess_audio, extract_ml_features
warnings.filterwarnings('ignore')
LABEL_MAP_INVERSE = {
0: 'Machine 1_Normal', 1: 'Machine 1_Abnormal',
2: 'Machine 2_Normal', 3: 'Machine 2_Abnormal',
4: 'Machine 3_Normal', 5: 'Machine 3_Abnormal'
}
# ==========================================
# LOAD MODEL
# ==========================================
model = None
model_error = ""
try:
model = joblib.load("ml_XGBoost_v1_immed.joblib")
print("XGBoost model loaded successfully.")
except Exception as e:
model_error = str(e)
print(f"Warning: Could not load model. Error: {model_error}")
# ==========================================
# PREDICTION
# ==========================================
def predict(audio_filepath):
if model is None:
return f"Model not loaded properly. Error: {model_error}"
if audio_filepath is None:
return "Please upload an audio file."
try:
# 1. Preprocess
y = preprocess_audio(audio_filepath)
# 2. Extract features
features = extract_ml_features(y)
# 3. Convert to array in the same column order the model expects
feature_names = sorted(features.keys())
X = np.array([[features[name] for name in feature_names]])
# 4. Predict
predicted_class = model.predict(X)[0]
predicted_label = LABEL_MAP_INVERSE.get(int(predicted_class), "Unknown")
# Get probabilities if available
if hasattr(model, 'predict_proba'):
proba = model.predict_proba(X)[0]
confidence = float(np.max(proba))
return f"Prediction: {predicted_label} (Confidence: {confidence:.2f})"
else:
return f"Prediction: {predicted_label}"
except Exception as e:
return f"Error processing file: {str(e)}"
# ==========================================
# GRADIO UI
# ==========================================
iface = gr.Interface(
fn=predict,
inputs=gr.Audio(type="filepath", label="Upload Machine Audio"),
outputs="text",
title="Machine Listener Diagnosis",
description="Upload a sound from a machine to predict whether it is Normal or Abnormal."
)
if __name__ == "__main__":
iface.launch()