fix language detection

main.py

@@ -17,21 +17,19 @@ import tempfile
 
 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)
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
-kit = None
-language_detector = None
+audio_processor = None
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    global kit, language_detector
-    kit = WhisperLiveKit()
-    language_detector = LanguageDetector(model_name="large")
+    global audio_processor
+    kit = WhisperLiveKit(args=args)
+    audio_processor = AudioProcessor()
     yield
 
 app = FastAPI(lifespan=lifespan)
@@ -47,8 +45,6 @@ app.add_middleware(
 # Mount static files
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-
-
 @app.get("/")
 async def read_root():
     return FileResponse("static/index.html")
@@ -66,9 +62,16 @@ async def detect_language(file: UploadFile = File(...)):
             contents = await file.read()
             temp_file.write(contents)
         
-        # Use the language detector with the saved file
-        if language_detector:
-            detected_lang, confidence, probs = language_detector.detect_language_from_file(file_path)
+        # Use the audio processor for language detection
+        if audio_processor:
+            # Load audio using librosa
+            audio, sr = librosa.load(file_path, sr=16000)
+            
+            # Convert to format expected by Whisper
+            audio = (audio * 32768).astype(np.int16)
+            
+            # Detect language
+            detected_lang, confidence, probs = audio_processor.detect_language(audio)
             
             # Clean up - remove the temporary file
             os.remove(file_path)
@@ -80,7 +83,7 @@ async def detect_language(file: UploadFile = File(...)):
             })
         else:
             return JSONResponse(
-                {"error": "Language detector not initialized"},
+                {"error": "Audio processor not initialized"},
                 status_code=500
             )
             
@@ -127,14 +130,15 @@ async def handle_websocket_results(websocket, results_generator):
 @app.websocket("/asr")
 async def websocket_endpoint(websocket: WebSocket):
     logger.info("New WebSocket connection request")
-    audio_processor = None
     websocket_task = None
 
     try:
         await websocket.accept()
         logger.info("WebSocket connection accepted")
         
-        audio_processor = AudioProcessor()
+        if not audio_processor:
+            raise RuntimeError("Audio processor not initialized")
+            
         results_generator = await audio_processor.create_tasks()
         websocket_task = asyncio.create_task(handle_websocket_results(websocket, results_generator))
 
@@ -155,8 +159,6 @@ async def websocket_endpoint(websocket: WebSocket):
         logger.error(f"Error in WebSocket endpoint: {e}")
         logger.error(f"Traceback: {traceback.format_exc()}")
     finally:
-        if audio_processor:
-            await audio_processor.cleanup()
         if websocket_task:
             websocket_task.cancel()
             try:

whisper_streaming_custom/backends.py

@@ -89,6 +89,42 @@ class WhisperTimestampedASR(ASRBase):
     def set_translate_task(self):
         self.transcribe_kargs["task"] = "translate"
 
+    def detect_language(self, audio):
+        import whisper
+        """
+        Detect the language of the audio using Whisper's language detection.
+        
+        Args:
+            audio (np.ndarray): Audio data as numpy array
+            
+        Returns:
+            tuple: (detected_language, confidence, probabilities)
+                - detected_language (str): The detected language code
+                - confidence (float): Confidence score for the detected language
+                - probabilities (dict): Dictionary of language probabilities
+        """
+        try:
+            # Ensure audio is in the correct format
+            if not isinstance(audio, np.ndarray):
+                audio = np.array(audio)
+            
+            # Pad or trim audio to the correct length
+            audio = whisper.pad_or_trim(audio)
+            
+            # Create mel spectrogram with correct dimensions
+            mel = whisper.log_mel_spectrogram(audio, n_mels=128).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, probs
+            
+        except Exception as e:
+            logger.error(f"Error in language detection: {e}")
+            raise
+
 
 class FasterWhisperASR(ASRBase):
     """Uses faster-whisper as the backend."""
@@ -147,6 +183,41 @@ class FasterWhisperASR(ASRBase):
     def set_translate_task(self):
         self.transcribe_kargs["task"] = "translate"
 
+    def detect_language(self, audio):
+        """
+        Detect the language of the audio using faster-whisper's language detection.
+        
+        Args:
+            audio (np.ndarray): Audio data as numpy array
+            
+        Returns:
+            tuple: (detected_language, confidence, probabilities)
+                - detected_language (str): The detected language code
+                - confidence (float): Confidence score for the detected language
+                - probabilities (dict): Dictionary of language probabilities
+        """
+        try:
+            # Ensure audio is in the correct format
+            if not isinstance(audio, np.ndarray):
+                audio = np.array(audio)
+            
+            # Use faster-whisper's detect_language method
+            language, language_probability, all_language_probs = self.model.detect_language(
+                audio=audio,
+                vad_filter=False,  # Disable VAD for language detection
+                language_detection_segments=1,  # Use single segment for detection
+                language_detection_threshold=0.5  # Default threshold
+            )
+            
+            # Convert list of tuples to dictionary for consistent return format
+            probs = {lang: prob for lang, prob in all_language_probs}
+            
+            return language, language_probability, probs
+            
+        except Exception as e:
+            logger.error(f"Error in language detection: {e}")
+            raise
+
 
 class MLXWhisper(ASRBase):
     """
@@ -225,6 +296,9 @@ class MLXWhisper(ASRBase):
 
     def set_translate_task(self):
         self.transcribe_kargs["task"] = "translate"
+    
+    def detect_language(self, audio):
+        raise NotImplementedError("MLX Whisper does not support language detection.")
 
 
 class OpenaiApiASR(ASRBase):
@@ -292,4 +366,7 @@ class OpenaiApiASR(ASRBase):
         self.use_vad_opt = True
 
     def set_translate_task(self):
-        self.task = "translate"
+        self.task = "translate"
+    
+    def detect_language(self, audio):
+        raise NotImplementedError("MLX Whisper does not support language detection.")