Spaces:
Sleeping
Sleeping
File size: 2,481 Bytes
feac291 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from flask import Flask, request, jsonify
import torch
import librosa
import numpy as np
from transformers import AutoModelForAudioClassification, Wav2Vec2FeatureExtractor
import os
app = Flask(__name__)
# Model setup
model_name = 'amiriparian/ExHuBERT'
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/hubert-base-ls960")
model = AutoModelForAudioClassification.from_pretrained(model_name, trust_remote_code=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# Labels for emotion mapping
labels = ['disgust', 'neutral', 'kind', 'anger', 'surprise', 'joy']
@app.route('/detect_scream', methods=['POST'])
def detect_scream():
try:
# Check if audio file is provided
if 'file' not in request.files:
return jsonify({'error': 'No audio file provided'}), 400
audio_file = request.files['file']
# Validate file type
if not audio_file.filename.endswith(('.wav', '.mp3')):
return jsonify({'error': 'Unsupported file format. Use WAV or MP3'}), 400
# Save audio file temporarily
temp_path = f"/tmp/{audio_file.filename}"
audio_file.save(temp_path)
# Load and preprocess audio
waveform, sr = librosa.load(temp_path, sr=16000)
inputs = feature_extractor(
waveform,
sampling_rate=16000,
padding="max_length",
max_length=48000,
return_tensors="pt"
)
inputs = inputs['input_values'].to(device)
# Perform inference
with torch.no_grad():
outputs = model(inputs).logits
probabilities = torch.nn.functional.softmax(outputs, dim=1)
confidence, predicted = torch.max(probabilities, 1)
# Get result
result = {
'label': labels[predicted.item()],
'confidence': float(confidence.item()),
'alert_level': 'High-Risk' if confidence.item() > 0.8 else ('Medium-Risk' if confidence.item() > 0.5 else 'None')
}
# Clean up temporary file
os.remove(temp_path)
return jsonify(result), 200
except Exception as e:
return jsonify({'error': str(e)}), 500
@app.route('/health', methods=['GET'])
def health_check():
return jsonify({'status': 'healthy'}), 200
if __name__ == '__main__':
app.run(host='0.0.0.0', port=7860) |