File size: 4,916 Bytes
9b6a9d7 e2a36e4 9b6a9d7 |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
from flask import Flask, request, jsonify, render_template
from transformers import pipeline
import torch
import torchaudio
from pydub import AudioSegment
import os
import io
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
from transformers import AutoModelForAudioClassification, AutoFeatureExtractor
app = Flask(__name__)
# Обновлённый словарь эмоций с surprise
emotion_map = {
'joy': '😊 Радость',
'neutral': '😐 Нейтрально',
'anger': '😠 Злость',
'sadness': '😢 Грусть',
'surprise': '😲 Удивление'
}
# Добавим ffmpeg и ffprobe в PATH
os.environ["PATH"] = r"C:\ffmpeg\bin;" + os.environ["PATH"]
# Модель для текста
text_classifier = pipeline(
"text-classification",
model="cointegrated/rubert-tiny2-cedr-emotion-detection",
config="cointegrated/rubert-tiny2-cedr-emotion-detection",
tokenizer="cointegrated/rubert-tiny2-cedr-emotion-detection",
framework="pt" # Явно указываем PyTorch
)
# Модель для аудио
# Аудио модель: универсальная, с поддержкой русского
audio_classifier = pipeline(
"audio-classification",
model="superb/hubert-large-superb-er",
device=0 if torch.cuda.is_available() else -1
)
@app.route("/")
def index():
return render_template("index.html")
@app.route("/analyze", methods=["POST"])
def analyze_text():
data = request.get_json()
text = data.get("text", "")
if not text.strip():
return jsonify({"error": "Пустой текст."}), 400
predictions = text_classifier(text)[0]
if not predictions:
return jsonify({"emotion": "❓ Не определено", "confidence": 0.0})
top_prediction = max(predictions, key=lambda x: x["score"])
label = top_prediction["label"]
confidence = round(top_prediction["score"], 2)
if confidence < 0.5:
return jsonify({"emotion": "🤔 Неуверенно", "confidence": confidence})
emotion = emotion_map.get(label, "❓ Неизвестно")
return jsonify({"emotion": emotion, "confidence": confidence})
from pydub.utils import which
# Проверка на наличие ffmpeg
if not which("ffmpeg"):
print("FFmpeg не найден в системе!")
@app.route('/analyze_audio', methods=['POST'])
def analyze_audio():
try:
if 'audio' not in request.files:
return jsonify({'error': 'Аудио файл обязателен'}), 400
audio_file = request.files['audio']
temp_path = "temp_audio.wav"
# Конвертация в формат для модели
audio = AudioSegment.from_file(io.BytesIO(audio_file.read()))
audio = audio.set_frame_rate(16000).set_channels(1)
audio.export(temp_path, format="wav", codec="pcm_s16le") # Явно указываем кодек
# Анализ
result = audio_classifier(temp_path)
os.remove(temp_path)
# Отладочный вывод
print("Raw model output:", result)
# Сопоставление с эмоциями
# Новый маппинг
emotion_mapping = {
'hap': 'happy',
'sad': 'sad',
'neu': 'neutral',
'ang': 'angry'
}
# Сбор эмоций
emotions = {v: 0.0 for v in emotion_mapping.values()}
for item in result:
label = item['label'].lower()
score = item['score']
if label in emotion_mapping:
key = emotion_mapping[label]
emotions[key] += score
# Нормализация
total = sum(emotions.values())
if total == 0:
return jsonify({'error': 'Не удалось распознать эмоции'}), 500
emotions = {k: round(v / total, 4) for k, v in emotions.items()}
dominant_emotion = max(emotions.items(), key=lambda x: x[1])
# Маппинг для ответа
response_map = {
'happy': '😊 Радость',
'sad': '😢 Грусть',
'angry': '😠 Злость',
'neutral': '😐 Нейтрально'
}
return jsonify({
'emotion': response_map.get(dominant_emotion[0], 'неизвестно'),
'confidence': dominant_emotion[1],
'details': emotions
})
except Exception as e:
print(f"Ошибка анализа аудио: {str(e)}")
return jsonify({'error': 'Ошибка обработки аудио'}), 500
if __name__ == "__main__":
app.run(debug=True)
if __name__ == "__main__":
app.run(debug=True)
|