Spaces:

JsonDONSON
/

Machine-Sound-Anomaly-Detector

Sleeping

Update app.py

39611d1 verified 7 months ago

2.1 kB

import gradio as gr
import torch
from model import MachineSoundCNN 
from config import NUM_CLASSES

# Import YOUR exact pipeline just like infer.py does!
from preprocessing import preprocess_audio
from features import audio_to_tensor

# 1. Define your exact classes
CLASSES = [
    "Machine 1 Normal", "Machine 1 Abnormal",
    "Machine 2 Normal", "Machine 2 Abnormal",
    "Machine 3 Normal", "Machine 3 Abnormal"
]

# 2. Load the Model 
device = torch.device('cpu')
model = MachineSoundCNN(num_classes=NUM_CLASSES)
model.load_state_dict(torch.load('best_model.pth', map_location=device, weights_only=True))
model.eval() 

# 3. Preprocessing & Inference Function
def predict_machine_sound(audio_path):
    if audio_path is None:
        return "Please upload an audio file."

    try:
        # Step A: Your exact Preprocessing (Resample -> Noise Reduce -> Silence -> Normalize)
        audio, sr = preprocess_audio(audio_path)

        # Step B: Your exact Feature Extraction (Mel Spec -> Padding -> Tensor)
        tensor = audio_to_tensor(audio, sr, augment=False)
        
        # Step C: Add batch dimension and move to CPU
        tensor = tensor.unsqueeze(0).to(device)

        # Step D: Forward Pass
        with torch.no_grad():
            output = model(tensor)
            probabilities = torch.nn.functional.softmax(output[0], dim=0)

        # Format output for Gradio
        result = {CLASSES[i]: float(probabilities[i]) for i in range(len(CLASSES))}
        return result
        
    except Exception as e:
        return {f"Error processing audio: {str(e)}": 1.0}

# 4. Build the Web Interface
interface = gr.Interface(
    fn=predict_machine_sound,
    inputs=gr.Audio(type="filepath", label="Upload Machine Audio (.wav)"),
    outputs=gr.Label(num_top_classes=6, label="CNN Prediction Confidence"),
    title="Industrial Machine Sound Anomaly Detector",
    description="Upload an audio clip of an industrial machine. The Custom CNN will analyze the audio using the exact training pipeline.",
    flagging_mode="never"
)

if __name__ == "__main__":
    interface.launch()

	import gradio as gr
	import torch
	from model import MachineSoundCNN
	from config import NUM_CLASSES

	# Import YOUR exact pipeline just like infer.py does!
	from preprocessing import preprocess_audio
	from features import audio_to_tensor

	# 1. Define your exact classes
	CLASSES = [
	"Machine 1 Normal", "Machine 1 Abnormal",
	"Machine 2 Normal", "Machine 2 Abnormal",
	"Machine 3 Normal", "Machine 3 Abnormal"
	]

	# 2. Load the Model
	device = torch.device('cpu')
	model = MachineSoundCNN(num_classes=NUM_CLASSES)
	model.load_state_dict(torch.load('best_model.pth', map_location=device, weights_only=True))
	model.eval()

	# 3. Preprocessing & Inference Function
	def predict_machine_sound(audio_path):
	if audio_path is None:
	return "Please upload an audio file."

	try:
	# Step A: Your exact Preprocessing (Resample -> Noise Reduce -> Silence -> Normalize)
	audio, sr = preprocess_audio(audio_path)

	# Step B: Your exact Feature Extraction (Mel Spec -> Padding -> Tensor)
	tensor = audio_to_tensor(audio, sr, augment=False)

	# Step C: Add batch dimension and move to CPU
	tensor = tensor.unsqueeze(0).to(device)

	# Step D: Forward Pass
	with torch.no_grad():
	output = model(tensor)
	probabilities = torch.nn.functional.softmax(output[0], dim=0)

	# Format output for Gradio
	result = {CLASSES[i]: float(probabilities[i]) for i in range(len(CLASSES))}
	return result

	except Exception as e:
	return {f"Error processing audio: {str(e)}": 1.0}

	# 4. Build the Web Interface
	interface = gr.Interface(
	fn=predict_machine_sound,
	inputs=gr.Audio(type="filepath", label="Upload Machine Audio (.wav)"),
	outputs=gr.Label(num_top_classes=6, label="CNN Prediction Confidence"),
	title="Industrial Machine Sound Anomaly Detector",
	description="Upload an audio clip of an industrial machine. The Custom CNN will analyze the audio using the exact training pipeline.",
	flagging_mode="never"
	)

	if __name__ == "__main__":
	interface.launch()