f

main.py

@@ -13,6 +13,7 @@ import logging
 import io
 from pydub import AudioSegment
 from typing import List
+import asyncio
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -46,17 +47,28 @@ model = joblib.load('models/xgb_test.pkl')
 class ConnectionManager:
     def __init__(self):
         self.active_connections: List[WebSocket] = []
+        self.audio_buffers = {}
 
     async def connect(self, websocket: WebSocket):
         await websocket.accept()
         self.active_connections.append(websocket)
+        self.audio_buffers[websocket] = b''  # Initialize buffer for each connection
 
     def disconnect(self, websocket: WebSocket):
         self.active_connections.remove(websocket)
+        del self.audio_buffers[websocket]  # Clean up buffer
 
-    async def send_message(self, message: str):
-        for connection in self.active_connections:
-            await connection.send_text(message)
+    async def send_message(self, websocket: WebSocket, message: str):
+        await websocket.send_text(message)
+
+    def add_to_buffer(self, websocket: WebSocket, data: bytes):
+        self.audio_buffers[websocket] += data  # Accumulate data in the buffer
+
+    def get_buffer(self, websocket: WebSocket) -> bytes:
+        return self.audio_buffers[websocket]
+
+    def clear_buffer(self, websocket: WebSocket):
+        self.audio_buffers[websocket] = b''  # Clear the buffer
 
 manager = ConnectionManager()
 
@@ -78,23 +90,14 @@ def extract_features(audio):
     combined_features = np.hstack([mfccs, chroma, contrast, centroid])
     return combined_features
 
-async def process_audio_data(audio_data):
+async def process_audio_data(websocket: WebSocket):
+    audio_data = manager.get_buffer(websocket)
     try:
-        logger.info(f"Audio data type: {type(audio_data)}")
-        logger.info(f"Raw audio length: {len(audio_data)}")
-        # Attempt to convert audio data from webm/ogg to wav format using pydub
+        # Convert audio data from webm/ogg to wav format using pydub
         audio_segment = AudioSegment.from_file(io.BytesIO(audio_data), format="webm")
-    except Exception as e:
-        logger.error(f"Failed to convert audio data using pydub: {e}")
-        return
-
-    try:
-        # Export the audio segment to wav format
         wav_io = io.BytesIO()
         audio_segment.export(wav_io, format="wav")
         wav_io.seek(0)
-
-        # Read the audio data
         audio, sr = sf.read(wav_io, dtype='float32')
     except Exception as e:
         logger.error(f"Failed to read audio data: {e}")
@@ -103,20 +106,14 @@ async def process_audio_data(audio_data):
     if audio.ndim > 1:  # If audio has more than one channel, average them
         audio = np.mean(audio, axis=1)
     
-    logger.info(f"The len of audio: {len(audio)}")
-    logger.info("Extracting features")
     features = extract_features(audio)
     features = features.reshape(1, -1)
-
-    try:
-        prediction = model.predict(features)
-        is_fake = prediction[0]
-        result = 'fake' if is_fake else 'real'
-    except Exception as e:
-        logger.error(f"Model prediction failed: {e}")
-        return
-
-    await manager.send_message(result)
+    prediction = model.predict(features)
+    is_fake = prediction[0]
+    result = 'fake' if is_fake else 'real'
+    
+    await manager.send_message(websocket, result)
+    manager.clear_buffer(websocket)
 
 @app.post("/start_detection")
 async def start_detection():
@@ -138,7 +135,8 @@ async def websocket_endpoint(websocket: WebSocket):
     try:
         while True:
             data = await websocket.receive_bytes()
-            await process_audio_data(data)
+            manager.add_to_buffer(websocket, data)
+            await process_audio_data(websocket)
     except WebSocketDisconnect:
         manager.disconnect(websocket)