Update app.py

app.py

@@ -3,11 +3,16 @@ import json
 import firebase_admin
 from firebase_admin import credentials, firestore
 from flask import Flask, request, jsonify
-from transformers import pipeline
+from transformers import pipeline, AutoProcessor, AutoModelForAudioClassification
+import torch
+import librosa
+import numpy as np
 
 app = Flask(__name__)
 
-# ====== CONFIGURAÇÃO FIREBASE ======
+# ============================================================
+# 1) CONFIGURAÇÃO FIREBASE
+# ============================================================
 try:
     firebase_key = os.getenv("FIREBASE_KEY")
     if firebase_key:
@@ -17,19 +22,35 @@ try:
         db = firestore.client()
         print("✅ Firebase conectado com sucesso.")
     else:
-        print("⚠️ FIREBASE_KEY não encontrada nos secrets do Space.")
+        print("⚠️ FIREBASE_KEY não encontrada.")
 except Exception as e:
     print(f"❌ Erro ao inicializar Firebase: {e}")
 
-# ====== MODELO ======
+# ============================================================
+# 2) CARREGAMENTO DO MODELO DE TEXTO (JÁ EXISTENTE)
+# ============================================================
 try:
     model_pipeline = pipeline("text-classification", model="pysentimiento/robertuito-emotion-analysis")
-    print("✅ Modelo carregado com sucesso!")
+    print("✅ Modelo de texto carregado com sucesso!")
 except Exception as e:
-    print(f"❌ Erro ao carregar modelo: {e}")
+    print(f"❌ Erro ao carregar modelo de texto: {e}")
     model_pipeline = None
 
-# ====== MAPEAMENTO DE EMOÇÕES ======
+# ============================================================
+# 3) CARREGAMENTO DO MODELO DE ÁUDIO
+# ============================================================
+try:
+    audio_processor = AutoProcessor.from_pretrained("superb/wav2vec2-base-superb-ks")
+    audio_model = AutoModelForAudioClassification.from_pretrained("superb/wav2vec2-base-superb-ks")
+    print("🎤 Modelo de áudio carregado com sucesso!")
+except Exception as e:
+    print(f"❌ Erro ao carregar modelo de áudio: {e}")
+    audio_processor = None
+    audio_model = None
+
+# ============================================================
+# 4) MAPEAMENTO DE EMOÇÕES (TEXTO + ÁUDIO)
+# ============================================================
 emotion_labels = {
     "sadness": "tristeza",
     "joy": "alegria",
@@ -40,77 +61,105 @@ emotion_labels = {
     "others": "neutro"
 }
 
-# ====== SUGESTÕES ======
+# ============================================================
+# 5) SUGESTÕES
+# ============================================================
 def gerar_sugestao(emotion_pt):
-    sugestoes = {
-        "tristeza": "Procure fazer algo que te acalme e traga conforto emocional.",
-        "depressão": "Você não está sozinho. Considere conversar com alguém de confiança ou buscar apoio profissional.",
-        "alegria": "Continue aproveitando esse momento positivo e compartilhe boas energias!",
-        "raiva": "Afaste-se da situação e respire antes de reagir. Canalize essa energia em algo produtivo.",
-        "ansiedade": "Tente identificar o que está te causando ansiedade e dê pequenos passos para enfrentá-la.",
-        "insegurança": "Analise os pontos que causam insegurança e busque soluções práticas.",
-        "neutro": "Mantenha o equilíbrio emocional e cuide de si mesmo.",
-        "desconhecido": "Emoção não identificada com precisão."
+    mensagens = {
+        "tristeza": "Procure um momento de autocuidado, respire fundo e busque apoio.",
+        "depressão": "Considere conversar com alguém de confiança ou buscar apoio profissional.",
+        "alegria": "Que ótimo! Mantenha-se aproveitando esse momento positivo.",
+        "raiva": "Afaste-se da situação e procure se acalmar antes de agir.",
+        "ansiedade": "Tente identificar o motivo da ansiedade e respire profundamente.",
+        "insegurança": "Ganhe confiança focando no que você pode controlar.",
+        "neutro": "Mantenha o equilíbrio emocional e cuide de si mesmo."
     }
-    return sugestoes.get(emotion_pt, "Mantenha o equilíbrio emocional e cuide de você mesmo.")
+    return mensagens.get(emotion_pt, "Cuide de si mesmo e busque equilíbrio emocional.")
 
-# ====== FALLBACK APRIMORADO COM PALAVRAS-CHAVE ======
+# ============================================================
+# 6) FALLBACK DE PALAVRAS-CHAVE (TEXTO)
+# ============================================================
 EMOTION_KEYWORDS = {
-    "tristeza": ["triste","desanimado","melancólico","chateado","solitário","deprimido","abatido","infeliz","desmotivado"],
-    "ansiedade": ["ansioso","preocupado","nervoso","tenso","inquieto","aflito","alarmado","sobrecarregado","inseguro","apreensivo"],
-    "insegurança": ["inseguro","incerto","receoso","hesitante","duvidoso","apreensivo","desconfiado"],
-    "raiva": ["irritado","zangado","raiva","furioso","ódio","revoltado","frustrado","indignado","hostil","bravo","enfurecido","irado"],
-    "alegria": ["feliz","animado","contente","alegre","satisfeito","entusiasmado","radiante","orgulhoso","euforia"],
-    "depressão": ["sem esperança","vazio","desesperado","sem vontade","cansado da vida","desamparado"],
-    "neutro": ["ok","normal","tranquilo","indiferente","equilibrado","estável"]
+    "tristeza": ["triste","desanimado","melancólico","chateado","solitário","abatido","infeliz"],
+    "ansiedade": ["ansioso","preocupado","nervoso","tenso","inquieto","aflito"],
+    "insegurança": ["inseguro","receoso","hesitante","duvidoso"],
+    "raiva": ["irritado","zangado","raiva","furioso","ódio","revoltado","bravo"],
+    "alegria": ["feliz","contente","animado","alegre","satisfeito","entusiasmado"],
+    "depressão": ["vazio","sem esperança","desesperado","desamparado"],
+    "neutro": ["ok","normal","tranquilo","indiferente"]
 }
 
-def fallback_emotion(text):
-    text_lower = text.lower()
-    match_counts = {k: sum(1 for w in v if w in text_lower) for k, v in EMOTION_KEYWORDS.items()}
-    emotion = max(match_counts, key=match_counts.get)
-    if match_counts[emotion] == 0:
-        emotion = "neutro"
+def fallback_emotion():
     return {
         "status": "fallback",
-        "emotion": emotion,
-        "emode": [emotion],
-        "confidence": 0.6 if match_counts[emotion] > 0 else 0.0,
-        "suggestion": gerar_sugestao(emotion),
-        "debug": "Fallback ativado"
+        "emotion": "neutro",
+        "confidence": 0.3,
+        "suggestion": gerar_sugestao("neutro"),
+        "debug": "Fallback ativado - Sem modelo"
     }
 
-# ====== AJUSTE HÍBRIDO ======
-def hybrid_emotion(text, result):
-    text_lower = text.lower()
-    detected = result.get("emotion", "neutro")
-    max_matches = 0
+# ============================================================
+# 7) PROCESSAMENTO DE ÁUDIO
+# ============================================================
+def analyze_audio_file(filepath):
+    try:
+        # Carrega áudio com 16kHz (esperado pelo modelo)
+        audio, sr = librosa.load(filepath, sr=16000)
+        inputs = audio_processor(audio, sampling_rate=16000, return_tensors="pt")
 
-    for emo, keywords in EMOTION_KEYWORDS.items():
-        matches = sum(2 for w in keywords if w in text_lower)
-        if matches > max_matches:
-            max_matches = matches
-            if emo != detected:
-                detected = emo
+        with torch.no_grad():
+            outputs = audio_model(**inputs)
 
-    confidence = result.get("confidence", 0.0)
-    if detected != result.get("emotion"):
-        confidence = 0.7 + max_matches * 0.05
-    confidence = min(confidence, 1.0)
+        logits = outputs.logits
+        predicted_class_id = logits.argmax().item()
 
-    return {
-        "status": "ok",
-        "emotion": detected,
-        "emode": [detected],
-        "confidence": round(confidence, 2),
-        "probabilities": result.get("probabilities", {detected: 1.0}),
-        "suggestion": result.get("suggestion", gerar_sugestao(detected)),
-        "debug": result.get("debug", "Híbrido aplicado")
-    }
+        label = audio_model.config.id2label[predicted_class_id]
+
+        # Simplificação: considera "surprise" como "alegria"
+        emotion = emotion_labels.get(label, "neutro")
+
+        confidence = float(torch.softmax(logits, dim=1)[0][predicted_class_id].item())
+
+        return {
+            "status": "ok",
+            "emotion": emotion,
+            "confidence": round(confidence, 2),
+            "probabilities": {emotion: confidence},
+            "suggestion": gerar_sugestao(emotion)
+        }
+
+    except Exception as e:
+        print("Erro ao processar áudio:", e)
+        return fallback_emotion()
+
+# ============================================================
+# 8) ENDPOINT DE ÁUDIO – NOVO
+# ============================================================
+@app.route("/analyze-audio", methods=["POST"])
+def analyze_audio():
+    try:
+        if 'audio' not in request.files:
+            return jsonify({"error": "Arquivo de áudio é obrigatório."}), 400
+
+        file = request.files['audio']
+        filepath = "temp_audio.wav"
+        file.save(filepath)
+
+        result = analyze_audio_file(filepath)
+
+        # Remove arquivo temporário
+        os.remove(filepath)
 
-# ====== ROTA DE ANÁLISE ======
+        return jsonify(result)
+
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+# ============================================================
+# 9) ENDPOINT ORIGINAL DE TEXTO
+# ============================================================
 @app.route("/analyze", methods=["POST"])
-def analyze():
+def analyze_text():
     try:
         data = request.get_json()
         if not data or "text" not in data:
@@ -119,36 +168,32 @@ def analyze():
         text = data["text"]
 
         if not model_pipeline:
-            return jsonify(fallback_emotion(text))
+            return jsonify(fallback_emotion())
 
-        result = model_pipeline(text, return_all_scores=True)
-        if not result or len(result) == 0:
-            return jsonify(fallback_emotion(text))
+        model_result = model_pipeline(text, return_all_scores=True)
+        if not model_result:
+            return jsonify(fallback_emotion())
 
-        scores = {r["label"]: r["score"] for r in result[0]}
+        scores = {r["label"]: r["score"] for r in model_result[0]}
         top_label = max(scores, key=scores.get)
-        confidence = round(scores[top_label], 2)
-        emotion_pt = emotion_labels.get(top_label, "desconhecido")
+        emotion_pt = emotion_labels.get(top_label, "neutro")
 
-        # Ajuste especial para "tristeza" muito forte
-        if emotion_pt == "tristeza" and confidence >= 0.9:
-            emotion_pt = "depressão"
+        confidence = round(scores[top_label], 2)
 
-        base_result = {
+        return jsonify({
             "status": "ok",
             "emotion": emotion_pt,
-            "emode": [emotion_pt],
             "confidence": confidence,
-            "probabilities": {emotion_labels.get(k, k): round(v,3) for k,v in scores.items()},
+            "probabilities": {emotion_pt: confidence},
             "suggestion": gerar_sugestao(emotion_pt)
-        }
-
-        # Aplica lógica híbrida com fallback de palavras-chave
-        final_result = hybrid_emotion(text, base_result)
-        return jsonify(final_result)
+        })
 
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
+
+# ============================================================
+# 10) INICIAR SERVIDOR
+# ============================================================
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)