Upload folder using huggingface_hub

asr_diarization/pipeline.py

@@ -44,7 +44,7 @@ class ASR_Diarization:
     def run_diarization(self, audio_path):
         diarization = self.diar_pipeline(audio_path)
         return [
-            {"start": t.start, "end": t.end, "speaker": spk}
+            {"segment_start": t.start, "segment_end": t.end, "speaker": spk}
             for t, _, spk in diarization.itertracks(yield_label=True)
         ]
 
@@ -54,8 +54,8 @@ class ASR_Diarization:
         speaker_segments = {}
 
         for seg in diar_json:
-            start, end, spk = seg["start"], seg["end"], seg["speaker"]
-            start_sample, end_sample = int(start * sr), int(end * sr)
+            segment_start, segment_end, spk = seg["segment_start"], seg["segment_end"], seg["speaker"]
+            start_sample, end_sample = int(segment_start * sr), int(segment_end * sr)
             chunk = audio[0, start_sample:end_sample].numpy()
 
             reduced = nr.reduce_noise(y=chunk, sr=sr)
@@ -74,8 +74,8 @@ class ASR_Diarization:
 
             seg_dict = {
                 "speaker": spk,
-                "segment_start": start,
-                "segment_end": end,
+                "segment_start": segment_start,
+                "segment_end": segment_end,
                 "tokens": tokens
             }
             merged_segments.append(seg_dict)
@@ -99,8 +99,8 @@ class ASR_Diarization:
             with open(rttm_path, "w") as f:
                 for seg in diar_json:
                     f.write(
-                        f"SPEAKER {base_name} 1 {seg['start']:.6f} "
-                        f"{seg['end']-seg['start']:.6f} <NA> <NA> "
+                        f"SPEAKER {base_name} 1 {seg['segment_start']:.6f} "
+                        f"{seg['segment_end']-seg['segment_start']:.6f} <NA> <NA> "
                         f"{seg['speaker']} <NA>\n"
                     )