Init sonogram_utility.py

Initialize sonogram_utility.py code

sonogram_utility.py

@@ -0,0 +1,96 @@
+import cv2
+import random
+import copy
+from pyannote.core import Annotation, Segment
+
+def colors(n):
+  '''
+  Creates a list size n of distinctive colors
+  '''
+  if n == 0:
+    return []
+  ret = []
+  h = int(random.random() * 180)
+  step = 180 / n
+  for i in range(n):
+    h += step
+    h = int(h) % 180
+    hsv = np.uint8([[[h,200,200]]])
+    bgr = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR)
+    ret.append((bgr[0][0][0].item()/255,bgr[0][0][1].item()/255,bgr[0][0][2].item()/255))
+  return ret
+
+def extendSpeakers(mySpeakerList, fileLabel = 'NONE', maximumSecondDifference = 1, minimumSecondDuration = 0):
+    '''
+    Assumes mySpeakerList is already split into Speaker/Audience 
+    '''
+    mySpeakerAnnotations = Annotation(uri=fileLabel)
+    newSpeakerList = [[],[]]
+    for i, speaker in enumerate(mySpeakerList):
+        speaker.sort()
+        lastEnd = -1
+        tempSection = None
+        for section in speaker:
+            if lastEnd == -1:
+                tempSection = copy.deepcopy(section)
+                lastEnd = section[0] + section[1]
+            else:
+                if section[0] - lastEnd <= maximumSecondDifference:
+                    tempSection = (tempSection[0],max(section[0] + section[1] - tempSection[0],tempSection[1]))
+                    lastEnd = tempSection[0] + tempSection[1]
+                else:
+                    if tempSection[1] >= minimumSecondDuration:
+                        newSpeakerList[i].append(tempSection)
+                        mySpeakerAnnotations[Segment(tempSection[0],lastEnd)] = i
+                    tempSection = copy.deepcopy(section)
+                    lastEnd = section[0] + section[1]
+        if tempSection is not None:
+            # Add the last section back in
+            if tempSection[1] >= minimumSecondDuration:
+                newSpeakerList[i].append(tempSection)
+                mySpeakerAnnotations[Segment(tempSection[0],lastEnd)] = i
+    return newSpeakerList,mySpeakerAnnotations
+
+def twoClassExtendAnnotation(myAnnotation,maximumSecondDifference = 1, minimumSecondDuration = 0):
+    lecturerID = None
+    lecturerLen = 0
+    
+    # Identify lecturer
+    for speakerName in myAnnotation.labels():
+        tempLen = len(myAnnotation.label_support(speakerName))
+        if tempLen > lecturerLen:
+            lecturerLen = tempLen
+            lecturerID = speakerName
+
+    tempSpeakerList = [[],[]]
+    # Recreate speakerList as [[lecturer labels],[audience labels]]
+    for speakerName in myAnnotation.labels():
+        if speakerName != lecturerID:
+            for segmentItem in myAnnotation.label_support(speakerName):
+                tempSpeakerList[1].append((segmentItem.start,segmentItem.duration))
+        else:
+            for segmentItem in myAnnotation.label_support(speakerName):
+                tempSpeakerList[0].append((segmentItem.start,segmentItem.duration))
+                
+    newList, newAnnotation = extendSpeakers(tempSpeakerList, fileLabel = myAnnotation.uri, maximumSecondDifference = maximumSecondDifference, minimumSecondDuration = minimumSecondDuration)
+
+    return newList, newAnnotation
+
+def loadAudioRTTM(sampleRTTM):
+    # Read in prediction data
+    # Data in list form, for convenient plotting
+    speakerList = []
+    # Data in Annotation form, for convenient error rate calculation
+    prediction = Annotation(uri=sampleRTTM)
+    with open(sampleRTTM, "r") as rttm:
+        for line in rttm:
+            speakerResult = line.split(' ')
+            index = int(speakerResult[7][-2:])
+            start = float(speakerResult[3])
+            end = start + float(speakerResult[4])
+            while len(speakerList) < index + 1:
+                speakerList.append([])
+            speakerList[index].append((float(speakerResult[3]),float(speakerResult[4])))
+            prediction[Segment(start,end)] = index
+
+    return speakerList, prediction