Add main.py

main.py

@@ -0,0 +1,83 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import numpy as np
+import librosa
+import re
+from unidecode import unidecode
+import base64
+
+app = FastAPI()
+
+class AudioBytesEncoded(BaseModel):
+    audio_bytes_encoded: str
+
+def get_model():
+    import faster_whisper
+    return faster_whisper.WhisperModel('tiny')
+
+model_faster = get_model()
+
+def CorregirErrores(texto):
+    for key, value in {'dinea': 'linea', 'dino': 'linea', 'tos': 'dos', 'dra': 'tra', 'una': 'uno', 'tes': 'tres', '1': 'uno', '2': 'dos', '3': 'tres'}.items():
+        texto = texto.replace(key, value)
+    return texto
+
+comando_base = ['linea', 'tra']
+comando_num = ['uno', 'dos', 'tres']
+
+pattern = '|'.join([re.escape(word) for word in comando_base + comando_num])
+
+nombre_clases = ['linea_uno', 'linea_dos', 'linea_tres', 'tra_uno', 'tra_dos', 'tra_tres']
+def predecir(audio):
+    resultado_final = None
+
+    completado = False
+    params1 = {'initial_prompt': 'Línea 1. Línea 2. Línea 3. Tra 1. Tra 2. Tra 3.',
+               'suppress_tokens': [],
+               'repetition_penalty': 2,
+               'no_speech_threshold': 0.1,
+               'log_prob_threshold': -0.1}
+    params2 = {'initial_prompt': [],
+               'suppress_tokens': [],
+               'repetition_penalty': 2,}
+    for params in (params1, params2):
+        for temp in [0, 1.0]:
+            resultado_original = model_faster.transcribe(audio, language='es', temperature=temp, **params)[0]
+            try:
+                resultado_original = next(resultado_original).text
+            except:
+                print('Falló la conversion.')
+                continue
+            print('Predicción:\t', resultado_original, end='\n')
+            resultado = unidecode(resultado_original.lower().strip())
+
+            resultado = CorregirErrores(resultado)
+            for resultado in resultado.split('.'):
+                matches = re.findall(pattern, resultado)
+
+                resultado_final = '_'.join(matches)
+                if resultado_final in nombre_clases:
+                    completado = True
+                    break
+            if completado:
+                break
+        if completado:
+            break
+
+    if resultado_final not in nombre_clases:
+        resultado = 'Comando no reconocido.'
+    return resultado
+
+@app.post("/predict/")
+async def predict(audio_bytes_encoded: AudioBytesEncoded):
+    try:
+        audio_bytes = base64.b64decode(audio_bytes_encoded.audio_bytes_encoded)
+        audio_np = np.frombuffer(audio_bytes, dtype=np.float32)
+        audio_np = librosa.util.normalize(audio_np)
+
+        prediction = predecir(audio_np)
+        return {"prediction": prediction}
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+