Update app.py

app.py

@@ -2,23 +2,19 @@ import joblib
 from transformers import AutoFeatureExtractor, Wav2Vec2Model
 import torch
 import librosa
-import numpy as np
-from sklearn.linear_model import LogisticRegression
-import gradio as gr
-import os
-import torch.nn.functional as F
 from scipy.special import expit
 import json
-
+import os
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+
+
 class CustomWav2Vec2Model(Wav2Vec2Model):
     def __init__(self, config):
         super().__init__(config)
         self.encoder.layers = self.encoder.layers[:9]
 
-truncated_model = CustomWav2Vec2Model.from_pretrained("facebook/wav2vec2-xls-r-2b")
 
 class HuggingFaceFeatureExtractor:
     def __init__(self, model, feature_extractor_name):
@@ -38,10 +34,14 @@ class HuggingFaceFeatureExtractor:
         inputs = {k: v.to(self.device) for k, v in inputs.items()}
         with torch.no_grad():
             outputs = self.model(**inputs, output_hidden_states=True)
-        return outputs.hidden_states[9]
+        return outputs.hidden_states[9] 
 
+
+truncated_model = CustomWav2Vec2Model.from_pretrained(r"C:\Users\david\PycharmProjects\David2\model\wav2vec2-xls-r-2b_truncated")
 FEATURE_EXTRACTOR = HuggingFaceFeatureExtractor(truncated_model, "facebook/wav2vec2-xls-r-2b")
-classifier,scaler, thresh = joblib.load('logreg_margin_pruning_ALL_with_scaler+threshold.joblib')
+classifier, scaler, thresh = joblib.load(r'C:\Users\david\PycharmProjects\David2\model\logreg_margin_pruning_ALL_with_scaler_threshold.joblib')
+
+
 
 def segment_audio(audio, sr, segment_duration):
     segment_samples = int(segment_duration * sr)
@@ -60,37 +60,33 @@ def process_audio(input_data, segment_duration=10):
         audio = audio[0]
     segments = segment_audio(audio, sr, segment_duration)
     segment_predictions = []
-    confidence_scores = []
-    eer_threshold = thresh - 5e-3  # small margin of error
-
+    confidence_scores_fake_sum = 0
+    fake_segments = 0
+    confidence_scores_real_sum = 0
+    real_segments = 0
+    eer_threshold = thresh - 5e-3 # small margin error due to feature extractor space differences
+    #print(eer_threshold)
     for idx, segment in enumerate(segments):
         features = FEATURE_EXTRACTOR(segment, sr)
-        features_avg = torch.mean(features, dim=1).cpu().numpy().reshape(1, -1)
+        features_avg = torch.mean(features, dim=1).cpu().numpy()
+        features_avg = features_avg.reshape(1, -1)
         decision_score = classifier.decision_function(features_avg)
         decision_score_scaled = scaler.transform(decision_score.reshape(-1, 1)).flatten()
         decision_value = decision_score_scaled[0]
         pred = 1 if decision_value >= eer_threshold else 0
-
-        if pred == 1:
-            confidence_percentage = expit(decision_score).item()
-        else:
+        if pred == 0:
             confidence_percentage = 1 - expit(decision_score).item()
-
+            confidence_scores_fake_sum +=confidence_percentage
+            fake_segments +=1
+        else:
+            confidence_percentage = expit(decision_score).item()
+            confidence_scores_real_sum +=confidence_percentage
+            real_segments +=1
         segment_predictions.append(pred)
-        confidence_scores.append(confidence_percentage)
-
     output_dict = {
-        "label": "real" if sum(segment_predictions) > (len(segment_predictions) / 2) else "fake",
-        "segments": [
-            {
-                "segment": idx + 1,
-                "prediction": "real" if pred == 1 else "fake",
-                "confidence": round(conf * 100, 2)
-            }
-            for idx, (pred, conf) in enumerate(zip(segment_predictions, confidence_scores))
-        ]
-    }
-
+    "label": "real" if sum(segment_predictions) > (len(segment_predictions) / 2) else "fake",
+    "confidence score:": f'{confidence_scores_real_sum/real_segments:.2f}' if sum(segment_predictions) > (len(segment_predictions) / 2) else f'{confidence_scores_fake_sum/fake_segments:.2f}'
+        }
     json_output = json.dumps(output_dict, indent=4)
     print(json_output)
     return json_output