Upload 2 files

ctcalign.py

@@ -0,0 +1,222 @@
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+import torch
+import numpy as np
+from dataclasses import dataclass
+
+
+
+#convert frame-numbers to timestamps in seconds
+# w2v2 step size is about 20ms, or 50 frames per second
+def f2s(fr):
+	return fr/50
+	
+# build labels dict from a processor where it is not directly accessible
+def get_processor_labels(processor,word_sep,max_labels=100):
+	ixs = sorted(list(range(max_labels)),reverse=True)
+	return {processor.tokenizer.decode(n) or word_sep:n for n in ixs}
+	
+#------------------------------------------
+# setup wav2vec2
+#------------------------------------------
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.random.manual_seed(0)
+max_labels = 100 # any reasonable number higher than vocab + extra + special tokens in any language used
+
+# important to know for CTC decoding - potentially language/model dependent
+model_word_separator = '|'
+model_blank_token = '[PAD]'
+
+is_MODEL_PATH="../models/LVL/wav2vec2-large-xlsr-53-icelandic-ep10-1000h"
+
+
+
+
+model = Wav2Vec2ForCTC.from_pretrained(is_MODEL_PATH).to(device)
+processor = Wav2Vec2Processor.from_pretrained(is_MODEL_PATH)
+labels_dict = get_processor_labels(processor,model_word_separator)
+inverse_dict = {v:k for k,v in labels_dict.items()}
+all_labels = tuple(labels_dict.keys())
+blank_id = labels_dict[model_blank_token]
+
+
+#------------------------------------------
+# forced alignment with ctc decoder
+#   based on implementation of
+#   https://pytorch.org/audio/main/tutorials/forced_alignment_tutorial.html
+#------------------------------------------
+
+
+# return the label class probability of each audio frame
+# wav is the wav data already read in, NOT the file path.
+def get_frame_probs(wav):
+    with torch.inference_mode(): # similar to with torch.no_grad():
+        input_values = processor(wav,sampling_rate=16000).input_values[0]
+        input_values = torch.tensor(input_values, device=device).unsqueeze(0)
+        emits = model(input_values).logits
+        emits = torch.log_softmax(emits, dim=-1)
+    return emits[0].cpu().detach()
+
+
+def get_trellis(emission, tokens, blank_id):
+    
+    num_frame = emission.size(0)
+    num_tokens = len(tokens)
+    trellis = torch.empty((num_frame + 1, num_tokens + 1))
+    trellis[0, 0] = 0
+    trellis[1:, 0] = torch.cumsum(emission[:, 0], 0) # len of this slice of trellis is len of audio frames)
+    trellis[0, -num_tokens:] = -float("inf") # len of this slice of trellis is len of transcript tokens
+    trellis[-num_tokens:, 0] = float("inf")
+    for t in range(num_frame):
+        trellis[t + 1, 1:] = torch.maximum(
+            # Score for staying at the same token
+            trellis[t, 1:] + emission[t, blank_id],
+            # Score for changing to the next token
+            trellis[t, :-1] + emission[t, tokens],
+        )
+    return trellis
+
+
+
+@dataclass
+class Point:
+    token_index: int
+    time_index: int
+    score: float
+    
+@dataclass
+class Segment:
+    label: str
+    start: int
+    end: int
+    score: float
+
+    @property
+    def mfaform(self):
+        return f"{f2s(self.start)},{f2s(self.end)},{self.label}"
+
+    @property
+    def length(self):
+        return self.end - self.start
+    
+    
+    
+def backtrack(trellis, emission, tokens, blank_id):
+    # Note:
+    # j and t are indices for trellis, which has extra dimensions
+    # for time and tokens at the beginning.
+    # When referring to time frame index `T` in trellis,
+    # the corresponding index in emission is `T-1`.
+    # Similarly, when referring to token index `J` in trellis,
+    # the corresponding index in transcript is `J-1`.
+    j = trellis.size(1) - 1
+    t_start = torch.argmax(trellis[:, j]).item()
+    
+    path = []
+    for t in range(t_start, 0, -1):
+        # 1. Figure out if the current position was stay or change
+        # `emission[J-1]` is the emission at time frame `J` of trellis dimension.
+        # Score for token staying the same from time frame J-1 to T.
+        stayed = trellis[t - 1, j] + emission[t - 1, blank_id]
+        # Score for token changing from C-1 at T-1 to J at T.
+        changed = trellis[t - 1, j - 1] + emission[t - 1, tokens[j - 1]]
+
+        # 2. Store the path with frame-wise probability.
+        prob = emission[t - 1, tokens[j - 1] if changed > stayed else 0].exp().item()
+        # Return token index and time index in non-trellis coordinate.
+        path.append(Point(j - 1, t - 1, prob))
+
+        # 3. Update the token
+        if changed > stayed:
+            j -= 1
+            if j == 0:
+                break
+    else:
+        raise ValueError("Failed to align")
+    return path[::-1]
+
+
+def merge_repeats(path,transcript):
+    i1, i2 = 0, 0
+    segments = []
+    while i1 < len(path):
+        while i2 < len(path) and path[i1].token_index == path[i2].token_index: # while both path steps point to the same token index
+            i2 += 1
+        score = sum(path[k].score for k in range(i1, i2)) / (i2 - i1)
+        segments.append( # when i2 finally switches to a different token,
+            Segment(
+                transcript[path[i1].token_index],# to the list of segments, append the token from i1
+                path[i1].time_index, # time of the first path-point of that token
+                path[i2 - 1].time_index + 1, # time of the final path-point for that token.
+                score,
+            )
+        )
+        i1 = i2
+    return segments
+
+
+
+def merge_words(segments, separator):
+    words = []
+    i1, i2 = 0, 0
+    while i1 < len(segments):
+        if i2 >= len(segments) or segments[i2].label == separator:
+            if i1 != i2:
+                segs = segments[i1:i2]
+                word = "".join([seg.label for seg in segs])
+                score = sum(seg.score * seg.length for seg in segs) / sum(seg.length for seg in segs)
+                words.append(Segment(word, segments[i1].start, segments[i2 - 1].end, score))
+            i1 = i2 + 1
+            i2 = i1
+        else:
+            i2 += 1
+    return words
+
+
+
+#------------------------------------------
+# handle etc.
+#------------------------------------------
+
+
+# generate mfa format for character (phone) and word alignments
+# skip the word separator as it is not a phone
+def mfalike(chars,wds,wsep):
+	hed = ['Begin,End,Label,Type,Speaker\n']
+	wlines = [f'{w.mfaform},words,000\n' for w in wds]
+	slines = [f'{ch.mfaform},phones,000\n' for ch in chars if ch.label != wsep]
+	return (''.join(hed+wlines+slines))
+	
+# generate basic exportable list format for character OR word alignments
+# skip the word separator as it is not a phone
+def basic(segs,wsep="|"):
+    return [[s.label,f2s(s.start),f2s(s.end)] for s in segs if s.label != wsep]
+	
+
+# needs pad labels added to correctly time first segment
+# and therefore add word sep character as placeholder in transcript
+def prep_transcript(xcp):
+    xcp = xcp.replace(' ',model_word_separator)
+    label_ids = [labels_dict[c] for c in xcp]
+    label_ids = [blank_id] + label_ids + [blank_id]
+    xcp = f'{model_word_separator}{xcp}{model_word_separator}'
+    return xcp,label_ids
+	
+	
+def align(wav_data,transcript):
+	norm_transcript,rec_label_ids = prep_transcript(transcript)
+	emit = get_frame_probs(wav_data)
+	trellis = get_trellis(emit, rec_label_ids, blank_id)
+	path = backtrack(trellis, emit, rec_label_ids, blank_id)
+	
+	segments = merge_repeats(path,norm_transcript)
+	words = merge_words(segments, model_word_separator)
+	
+	#segments = [s for s in segments if s[0] != model_word_separator]
+	#return mfalike(segments,words,model_word_separator)
+	return basic(words,model_word_separator), basic(segments,model_word_separator)
+
+
+
+
+

graph.py

@@ -0,0 +1,117 @@
+import numpy as np
+import soundfile as sf
+from scipy import signal
+import librosa
+import subprocess
+import matplotlib.pyplot as plt
+import ctcalign
+
+
+
+def readwav(wav_path):
+    wav, sr = sf.read(wav_path, dtype=np.float32)
+    if len(wav.shape) == 2:
+        wav = wav.mean(1)
+    if sr != 16000:
+        wlen = int(wav.shape[0] / sr * 16000)
+        wav = signal.resample(wav, wlen)
+    return wav
+
+
+def normalise_transcript(xcp):
+    xcp = xcp.lower()
+    while '  ' in xcp:
+        xcp = xcp.replace('  ', ' ')
+    return xcp
+
+
+
+def get_pitch_tracks(wav_path):
+  f0_data = subprocess.run(["REAPER/build/reaper", "-i", wav_path, '-a']).stdout
+  #with open('tmp.f0','r') as handle:
+  f0_data = f0_data.split('EST_Header_End\n')[1].splitlines()
+  print(f0_data) #!!!!!!!!!!!!!!!!!!!!!
+  f0_data = [l.split(' ') for l in f0_data]
+  f0_data = [ [float(t), float(f)] for t,v,f in f0_data if v=='1']
+  f0_data = [[t,f0] for t,prob,f0 in f0_data if prob==1.0]
+  return f0_data
+
+
+
+
+
+
+
+# transcript could be from a corpus with the wav file,
+# input by the user,
+# or from a previous speech recognition process
+def align_and_graph(wav_path, transcript):
+
+    # fetch data
+    #f0_data = get_pitch_tracks(wav_path)
+    speech = readwav(wav_path)
+    w_align, seg_align = ctcalign.align(speech,normalise_transcript(transcript))
+    
+    
+    # set up the graph shape
+    rec_start = w_align[0][1]
+    rec_end = w_align[-1][2]
+    
+    f0_data = get_pitch_tracks(wav_path)
+    if f0_data:
+        f_max = max([f0 for t,f0 in f0_data]) + 50
+    else:
+        f_max = 400
+
+
+    fig, axes1 = plt.subplots(figsize=(15,5))
+    plt.xlim([rec_start, rec_end])
+    axes1.set_ylim([0.0, f_max])
+    axes1.get_xaxis().set_visible(False)
+    
+    # draw word boundaries
+    for w,s,e in w_align:
+        plt.vlines(s,0,f_max,linewidth=0.5,color='black')
+        plt.vlines(e,0,f_max,linewidth=0.5,color='black')
+        plt.text( (s+e)/2 - (len(w)*.01), f_max+15, w, fontsize=15)
+        
+    # draw phone/char boundaries
+    for p,s,e in seg_align:
+        plt.vlines(s,0,f_max,linewidth=0.3,color='cadetblue',linestyle=(0,(10,4)))
+        plt.vlines(e,0,f_max,linewidth=0.3,color='cadetblue',linestyle=(0,(10,4)))
+        plt.text( (s+e)/2 - (len(p)*.01), -30, p, fontsize=15, color='teal')
+    
+    
+    f0c = "blue"
+    axes1.scatter([t for t,f0 in f0_data], [f0 for t,f0 in f0_data], color=f0c)
+
+    
+    
+    w, sr = librosa.load(wav_path)
+    fr_l = 2048 # librosa default
+    h_l = 512 # default
+    rmse = librosa.feature.rms(y=w, frame_length = fr_l, hop_length = h_l)
+    rmse = rmse[0]
+
+
+    # show rms energy
+    axes2 = axes1.twinx()
+    axes2.set_ylim([0.0, 0.5])
+    rms_xval = [(h_l*i)/sr for i in range(len(rmse))]
+    axes2.plot(rms_xval,rmse,color='peachpuff',linewidth=3.5)
+
+    
+    # label the graph
+    axes1.set_ylabel("Pitch (F0, Hz)", fontsize=14, color="blue")
+    axes2.set_ylabel("RMS energy", fontsize=14,color="coral")
+    #plt.title(f'Recording {file_id} (L1 {language_dict[file_id]})', fontsize=15)
+    #plt.show()
+    
+    return fig
+    
+    #plt.close('all')
+
+
+# uppboðssøla bussleiðini viðmerkingar upprunaligur
+
+