Update app.py

app.py

@@ -5,7 +5,7 @@ from pydub.silence import detect_nonsilent
 from pydub import AudioSegment
  
 
-def get_labels(audio_fp, min_speech_duration_ms, min_silence_duration_ms, auto_merge, uppper_merge_threshold, max_segment_length):
+def get_labels(audio_fp, min_speech_duration_ms, min_silence_duration_ms, auto_merge, uppper_merge_threshold, max_segment_length, end_extension, extend_small_segments):
     audio = AudioSegment.from_file(audio_fp).set_channels(1)
     speech_timestamps = detect_nonsilent(audio, min_silence_len=min_silence_duration_ms, silence_thresh=-40)
     speech_timestamps = list(filter(lambda x: x[1]-x[0] > min_speech_duration_ms, speech_timestamps))
@@ -44,6 +44,16 @@ def get_labels(audio_fp, min_speech_duration_ms, min_silence_duration_ms, auto_m
         else:
             new_labels.append(list(labels[i]))
 
+    if extend_small_segments:
+        for i, nl in enumerate(new_labels[:-1]):
+            if nl[1] - nl[0] <= 1.2 and nl[0] + 1.2 < new_labels[i+1][0]:
+                nl[1] = nl[0] + 1.2
+
+    if end_extension:
+        for i, nl in enumerate(new_labels[:-1]):
+            if nl[0] + end_extension < new_labels[i+1][0]:
+                nl[1] = nl[0] + end_extension
+
     translate_labels = list(map(lambda x: f"{x[0]}\t{x[1]}\t{x[2]}", new_labels))
 
     filename_path = f"{fn}_translate_label.txt"
@@ -57,11 +67,13 @@ interface = gr.Interface(
     get_labels,
     [
         gr.Audio(type="filepath", label="Audio file"),
-        gr.Number(label="min_speech_duration_ms", value=80, info="default (80)"), 
-        gr.Number(label="min_silence_duration_ms", value=40, info="default (100)"),
+        gr.Number(label="min_speech_duration_ms", value=40, info="default (40)"), 
+        gr.Number(label="min_silence_duration_ms", value=40, info="default (40)"),
         gr.Checkbox(label="Auto merge", value=True),
         gr.Textbox(label="Gap max threshold value (seconds)", value=0.350),
-        gr.Number(label="Approx Max Segment Length", value=7)
+        gr.Number(label="Approx Max Segment Length", value=7),
+        gr.Number(label="Extend end by (seconds)", value=0),
+        gr.Checkbox(label="Extend small segments (minimum 1.2 seconds)", value=False)
     ],
     [
         gr.File(label="VAD Labels"),