feat: try to add language detector

language_detector.py

@@ -0,0 +1,84 @@
+import whisper
+import numpy as np
+import logging
+import io
+import librosa
+
+logger = logging.getLogger(__name__)
+
+class LanguageDetector:
+    def __init__(self, model_name="tiny"):
+        """
+        Initialize the language detector with a Whisper model.
+        
+        Args:
+            model_name (str): Name of the Whisper model to use. Default is "tiny" which is sufficient for language detection.
+        """
+        self.model = whisper.load_model(model_name)
+        logger.info(f"Loaded Whisper model {model_name} for language detection")
+
+    def detect_language_from_file(self, audio_file_path):
+        """
+        Detect language from an audio file.
+        
+        Args:
+            audio_file_path (str): Path to the audio file
+            
+        Returns:
+            str: Detected language code (e.g., "en", "fr", etc.)
+            float: Confidence score
+        """
+        try:
+            # Load and preprocess audio
+            audio = whisper.load_audio(audio_file_path)
+            audio = whisper.pad_or_trim(audio)
+            
+            # Make log-Mel spectrogram
+            mel = whisper.log_mel_spectrogram(audio).to(self.model.device)
+            
+            # Detect language
+            _, probs = self.model.detect_language(mel)
+            detected_lang = max(probs, key=probs.get)
+            confidence = probs[detected_lang]
+            
+            return detected_lang, confidence
+            
+        except Exception as e:
+            logger.error(f"Error in language detection: {e}")
+            raise
+
+    def detect_language_from_bytes(self, audio_bytes):
+        """
+        Detect language from audio bytes.
+        
+        Args:
+            audio_bytes (bytes): Audio data in bytes
+            
+        Returns:
+            str: Detected language code (e.g., "en", "fr", etc.)
+            float: Confidence score
+        """
+        try:
+            # Convert bytes to numpy array using librosa
+            audio_data = io.BytesIO(audio_bytes)
+            audio, sr = librosa.load(audio_data, sr=16000)
+            
+            # Convert to format expected by Whisper
+            audio = (audio * 32768).astype(np.int16)
+            
+            # Load and preprocess audio
+            audio = whisper.pad_or_trim(audio)
+            
+            # Make log-Mel spectrogram
+            mel = whisper.log_mel_spectrogram(audio).to(self.model.device)
+            
+            # Detect language
+            _, probs = self.model.detect_language(mel)
+            detected_lang = max(probs, key=probs.get)
+            confidence = probs[detected_lang]
+            
+            return detected_lang, confidence
+            
+        except Exception as e:
+            logger.error(f"Error in language detection: {e}")
+            raise 

main.py

@@ -1,5 +1,5 @@
 from contextlib import asynccontextmanager
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi.staticfiles import StaticFiles
@@ -10,9 +10,13 @@ import os
 import traceback
 import argparse
 import uvicorn
+import numpy as np
+import librosa
+import io
 
 from core import WhisperLiveKit
 from audio_processor import AudioProcessor
+from language_detector import LanguageDetector
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logging.getLogger().setLevel(logging.WARNING)
@@ -20,11 +24,13 @@ logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 kit = None
+language_detector = None
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    global kit
+    global kit, language_detector
     kit = WhisperLiveKit()
+    language_detector = LanguageDetector(model_name="tiny")
     yield
 
 app = FastAPI(lifespan=lifespan)
@@ -50,6 +56,34 @@ async def read_root():
 async def health_check():
     return JSONResponse({"status": "healthy"})
 
+@app.post("/detect-language")
+async def detect_language(file: UploadFile = File(...)):
+    try:
+        # Read the audio file
+        contents = await file.read()
+        
+        # Use the language detector
+        if language_detector:
+            detected_lang, confidence = language_detector.detect_language_from_bytes(contents)
+            
+            return JSONResponse({
+                "language": detected_lang,
+                "confidence": float(confidence)
+            })
+        else:
+            return JSONResponse(
+                {"error": "Language detector not initialized"},
+                status_code=500
+            )
+            
+    except Exception as e:
+        logger.error(f"Error in language detection: {e}")
+        logger.error(f"Traceback: {traceback.format_exc()}")
+        return JSONResponse(
+            {"error": str(e)},
+            status_code=500
+        )
+
 async def handle_websocket_results(websocket, results_generator):
     """Consumes results from the audio processor and sends them via WebSocket."""
     try: