Styling

app.py

@@ -5,12 +5,25 @@ from transformers import AutoFeatureExtractor, AutoModelForAudioXVector
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-OUTPUT = """
-    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" integrity="sha256-YvdLHPgkqJ8DVUxjjnGVlMMJtNimJ6dYkowFFvp4kKs=" crossorigin="anonymous">
+STYLE = """
+<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" integrity="sha256-YvdLHPgkqJ8DVUxjjnGVlMMJtNimJ6dYkowFFvp4kKs=" crossorigin="anonymous">
+"""
+OUTPUT_OK = STYLE + """
+    <div class="container">
+        <div class="row"><h1 style="text-align: center">The speakers are</h1></div>
+        <div class="row"><h1 class="display-1 text-success" style="text-align: center">{:.1f}%</h1></div>
+        <div class="row"><h1 style="text-align: center">similar</h1></div>
+        <div class="row"><h1 class="text-success" style="text-align: center">Welcome, human!</h1></div>
+        <div class="row"><small style="text-align: center">(You must get 89% or more to be considered the same person)</small><div class="row">
+    </div>
+"""
+OUTPUT_FAIL = STYLE + """
     <div class="container">
         <div class="row"><h1 style="text-align: center">The speakers are</h1></div>
-        <div class="row"><h1 class="display-1" style="text-align: center">{:.1f}%</h1></div>
+        <div class="row"><h1 class="display-1 text-danger" style="text-align: center">{:.1f}%</h1></div>
         <div class="row"><h1 style="text-align: center">similar</h1></div>
+        <div class="row"><h1 class="text-danger" style="text-align: center">You shall not pass!</h1></div>
+        <div class="row"><small style="text-align: center">(You must get 89% or more to be considered the same person)</small><div class="row">
     </div>
 """
 
@@ -21,7 +34,9 @@ EFFECTS = [
     ["silence", "1", "0.1", "0.1%", "-1", "0.1", "0.1%"],
 ]
 
-model_name = "anton-l/unispeech-sat-base-plus-sv"
+THRESHOLD = 0.89
+
+model_name = "microsoft/unispeech-sat-base-plus-sv"
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 model = AutoModelForAudioXVector.from_pretrained(model_name).to(device)
 cosine_sim = torch.nn.CosineSimilarity(dim=-1)
@@ -33,9 +48,10 @@ def similarity_fn(mic_path1, file_path1, mic_path2, file_path2):
 
     wav1, _ = apply_effects_file(mic_path1 if mic_path1 else file_path1, EFFECTS)
     wav2, _ = apply_effects_file(mic_path2 if mic_path2 else file_path2, EFFECTS)
+    print(wav1.shape, wav2.shape)
 
-    input1 = feature_extractor(wav1.squeeze(0), return_tensors="pt").input_values.to(device)
-    input2 = feature_extractor(wav2.squeeze(0), return_tensors="pt").input_values.to(device)
+    input1 = feature_extractor(wav1.squeeze(0), return_tensors="pt", sampling_rate=16000).input_values.to(device)
+    input2 = feature_extractor(wav2.squeeze(0), return_tensors="pt", sampling_rate=16000).input_values.to(device)
 
     with torch.no_grad():
         emb1 = model(input1).embeddings
@@ -44,7 +60,12 @@ def similarity_fn(mic_path1, file_path1, mic_path2, file_path2):
     emb2 = torch.nn.functional.normalize(emb2, dim=-1).cpu()
     similarity = cosine_sim(emb1, emb2).numpy()[0]
 
-    return OUTPUT.format(similarity * 100)
+    if similarity >= THRESHOLD:
+        output = OUTPUT_OK.format(similarity * 100)
+    else:
+        output = OUTPUT_FAIL.format(similarity * 100)
+
+    return output
 
 
 inputs = [
@@ -63,7 +84,8 @@ description = (
 article = (
     "<p style='text-align: center'>"
     "<a href='https://huggingface.co/microsoft/unispeech-sat-large' target='_blank'>🎙️ Learn more about UniSpeech-SAT</a> | "
-    "<a href='https://arxiv.org/abs/2110.05752' target='_blank'>📚 Article on ArXiv</a>"
+    "<a href='https://arxiv.org/abs/2110.05752' target='_blank'>📚 UniSpeech-SAT paper</a> | "
+    "<a href='https://www.danielpovey.com/files/2018_icassp_xvectors.pdf' target='_blank'>📚 X-Vector paper</a>"
     "</p>"
 )
 
@@ -71,7 +93,7 @@ interface = gr.Interface(
     fn=similarity_fn,
     inputs=inputs,
     outputs=output,
-    title="Speaker Verification with UniSpeech-SAT",
+    title="Speaker Verification with UniSpeech-SAT + X-Vectors",
     description=description,
     article=article,
     layout="horizontal",