Update voice_analysis.py

voice_analysis.py

@@ -1,40 +1,48 @@
-import torch
 import numpy as np
-from speechbrain.pretrained import EncoderClassifier
-from pydub import AudioSegment
-from sklearn.cluster import DBSCAN
 import librosa
+from sklearn.cluster import DBSCAN
+from pydub import AudioSegment
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-classifier = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb", savedir="pretrained_models/spkrec-ecapa-voxceleb", run_opts={"device": device})
-
-def extract_voice_embedding(audio_segment):
-    signal = np.array(audio_segment.get_array_of_samples())
-    signal = signal.astype(np.float32) / 32768.0  # Normalize to [-1, 1]
-    embedding = classifier.encode_batch(torch.tensor(signal).unsqueeze(0))
-    return embedding.squeeze().cpu().numpy()
-
-def process_audio(audio_path, segment_duration=1000):
-    audio = AudioSegment.from_file(audio_path)
-    segments = [audio[i:i+segment_duration] for i in range(0, len(audio), segment_duration)]
-    embeddings = [extract_voice_embedding(segment) for segment in segments]
-    return embeddings
-
-def cluster_voices(embeddings):
-    if len(embeddings) < 2:
+def extract_voice_features(audio_path, segment_duration=1000):
+    # Load the audio file
+    y, sr = librosa.load(audio_path)
+    
+    # Extract MFCC features
+    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
+    
+    # Segment the MFCCs
+    segment_length = int(segment_duration * sr / 1000)
+    num_segments = len(y) // segment_length
+    
+    segments = []
+    for i in range(num_segments):
+        start = i * segment_length
+        end = start + segment_length
+        segment = mfccs[:, start:end]
+        segments.append(np.mean(segment, axis=1))
+    
+    return np.array(segments)
+
+def cluster_voices(features):
+    if len(features) < 2:
         print("Not enough voice segments for clustering. Assigning all to one cluster.")
-        return np.zeros(len(embeddings), dtype=int)
+        return np.zeros(len(features), dtype=int)
 
-    X = np.stack(embeddings)
-    dbscan = DBSCAN(eps=0.3, min_samples=5, metric='cosine')
-    clusters = dbscan.fit_predict(X)
+    dbscan = DBSCAN(eps=0.5, min_samples=5, metric='euclidean')
+    clusters = dbscan.fit_predict(features)
 
     if np.all(clusters == -1):
         print("DBSCAN assigned all to noise. Considering as one cluster.")
-        return np.zeros(len(embeddings), dtype=int)
+        return np.zeros(len(features), dtype=int)
 
     return clusters
 
-def get_most_frequent_voice(embeddings, clusters):
+def get_most_frequent_voice(features, clusters):
     largest_cluster = max(set(clusters), key=list(clusters).count)
-    return [emb for emb, cluster in zip(embeddings, clusters) if cluster == largest_cluster]
+    return features[clusters == largest_cluster]
+
+def process_audio(audio_path, segment_duration=1000):
+    features = extract_voice_features(audio_path, segment_duration)
+    clusters = cluster_voices(features)
+    most_frequent_voice = get_most_frequent_voice(features, clusters)
+    return most_frequent_voice, features, clusters