Update scraibe/diarisation.py

scraibe/diarisation.py

@@ -87,68 +87,94 @@ class Diariser:
 
         return out
 
-    @staticmethod
-    def format_diarization_output(dia: Annotation) -> dict:
-        """
-        Formats the raw diarization output into a more usable structure for this project.
+    # @staticmethod
+    # def format_diarization_output(dia: Annotation) -> dict:
+    #     """
+    #     Formats the raw diarization output into a more usable structure for this project.
 
-        Args:
-            dia: Raw diarization output.
+    #     Args:
+    #         dia: Raw diarization output.
 
-        Returns:
-            dict: A structured representation of the diarization, with speaker names
-                  as keys and a list of tuples representing segments as values.
-        """
+    #     Returns:
+    #         dict: A structured representation of the diarization, with speaker names
+    #               as keys and a list of tuples representing segments as values.
+    #     """
 
-        dia_list = list(dia.itertracks(yield_label=True))
-        diarization_output = {"speakers": [], "segments": []}
+    #     dia_list = list(dia.itertracks(yield_label=True))
+    #     diarization_output = {"speakers": [], "segments": []}
 
-        normalized_output = []
-        index_start_speaker = 0
-        index_end_speaker = 0
-        current_speaker = str()
+    #     normalized_output = []
+    #     index_start_speaker = 0
+    #     index_end_speaker = 0
+    #     current_speaker = str()
 
-        ###
-        # Sometimes two consecutive speakers are the same
-        # This loop removes these duplicates
-        ###
+    #     ###
+    #     # Sometimes two consecutive speakers are the same
+    #     # This loop removes these duplicates
+    #     ###
 
-        if len(dia_list) == 1:
-            normalized_output.append([0, 0, dia_list[0][2]])
-        else:
+    #     if len(dia_list) == 1:
+    #         normalized_output.append([0, 0, dia_list[0][2]])
+    #     else:
+
+    #         for i, (_, _, speaker) in enumerate(dia_list):
 
-            for i, (_, _, speaker) in enumerate(dia_list):
+    #             if i == 0:
+    #                 current_speaker = speaker
 
-                if i == 0:
-                    current_speaker = speaker
+    #             if speaker != current_speaker:
 
-                if speaker != current_speaker:
+    #                 index_end_speaker = i - 1
 
-                    index_end_speaker = i - 1
+    #                 normalized_output.append([index_start_speaker,
+    #                                           index_end_speaker,
+    #                                           current_speaker])
 
-                    normalized_output.append([index_start_speaker,
-                                              index_end_speaker,
-                                              current_speaker])
+    #                 index_start_speaker = i
+    #                 current_speaker = speaker
 
-                    index_start_speaker = i
-                    current_speaker = speaker
+    #             if i == len(dia_list) - 1:
 
-                if i == len(dia_list) - 1:
+    #                 index_end_speaker = i
 
-                    index_end_speaker = i
+    #                 normalized_output.append([index_start_speaker,
+    #                                           index_end_speaker,
+    #                                           current_speaker])
 
-                    normalized_output.append([index_start_speaker,
-                                              index_end_speaker,
-                                              current_speaker])
+    #     for outp in normalized_output:
+    #         start = dia_list[outp[0]][0].start
+    #         end = dia_list[outp[1]][0].end
 
-        for outp in normalized_output:
-            start = dia_list[outp[0]][0].start
-            end = dia_list[outp[1]][0].end
+    #         diarization_output["segments"].append([start, end])
+    #         diarization_output["speakers"].append(outp[2])
+    #     return diarization_output
+
+    @staticmethod
+    def format_diarization_output(dia: Annotation) -> dict:
+        """
+            Formats the raw diarization output into a more usable structure for this project,
+            without combining consecutive segments of the same speaker.
+
+            Args:
+            dia: Raw diarization output.
+
+            Returns:
+            dict: A structured representation of the diarization, with speaker names
+              as keys and a list of tuples representing segments as values.
+        """
+        dia_list = list(dia.itertracks(yield_label=True))
+        diarization_output = {"speakers": [], "segments": []}
+
+        for segment, _, speaker in dia_list:
+            start = segment.start
+            end = segment.end
 
             diarization_output["segments"].append([start, end])
-            diarization_output["speakers"].append(outp[2])
+            diarization_output["speakers"].append(speaker)
+    
         return diarization_output
 
+
     @staticmethod
     def _get_token():
         """