align.py

#!/usr/bin/env python3

"""
Command-line usage:
  python align.py [options] wave_file transcript_file output_file
  where options may include:
	-r sampling_rate -- override which sample rate model to use, one of 8000, 11025, and 16000
	-s start_time    -- start of portion of wavfile to align (in seconds, default 0)
	-e end_time      -- end of portion of wavfile to align (in seconds, default to end)

You can also import this file as a module and use the functions directly.
"""

import os
import sys
import getopt
import wave
import re
import subprocess
import unicodedata


VOWEL_PHONES = {
	'a', 'ae', 'ya', 'yae', 'eo', 'e', 'yeo', 'ye',
	'o', 'wa', 'wae', 'oe', 'yo', 'u', 'wo', 'we', 'wi',
	'yu', 'eu', 'yi', 'i'
}


def _is_vowel(phone):
	if not phone:
		return False
	low = phone.lower()
	if low in VOWEL_PHONES:
		return True
	return low[0] in {'a', 'e', 'i', 'o', 'u', 'w', 'y'}


def prep_wav(orig_wav, out_wav, sr_override, wave_start, wave_end):
	global sr_models

	# If we had previously generated out_wav and wanted to reuse it, we could early-return.
	# Currently disabled by design (kept for reference).
	if os.path.exists(out_wav) and False:
		f = wave.open(out_wav, 'r')
		SR = f.getframerate()
		f.close()
		print("Already re-sampled the wav file to " + str(SR))
		return SR

	f = wave.open(orig_wav, 'r')
	SR = f.getframerate()
	f.close()

	soxopts = ""
	if float(wave_start) != 0.0 or wave_end is not None:
		soxopts += " trim " + wave_start
		if wave_end is not None:
			soxopts += " " + str(float(wave_end) - float(wave_start))

	# Resample if needed (model SR mismatch or override), or if we need to trim.
	if (sr_models is not None and SR not in sr_models) or (sr_override is not None and SR != sr_override) or soxopts != "":
		# Default to 16000 Hz for better quality (was 11025)
		new_sr = 16000
		if sr_override is not None:
			new_sr = sr_override

		print("Resampling wav file from " + str(SR) + " to " + str(new_sr) + soxopts + "...")
		SR = new_sr
		# Correct sox syntax: sox input output [effects]
		os.system("sox \"" + orig_wav + "\" \"" + out_wav + "\" rate -v " + str(SR) + soxopts)
	else:
		# Already at the desired sample rate and no trimming required.
		os.system("cp -f " + orig_wav + " " + out_wav)

	return SR


def _read_text_any_encoding(path):
	"""Read text file trying common encodings and return a unicode string.
	Tries: utf-8, utf-16, utf-16le, utf-16be, cp949, euc-kr. Strips leading BOM if present."""
	encodings = [
		'utf-8', 'utf-16', 'utf-16le', 'utf-16be', 'cp949', 'euc-kr'
	]
	data = None
	with open(path, 'rb') as fb:
		data = fb.read()
	for enc in encodings:
		try:
			s = data.decode(enc)
			# strip BOM if any survived
			if s and s[0] == '\ufeff':
				s = s[1:]
			return s
		except Exception:
			continue
	# Fallback: best-effort utf-8 with replacement
	try:
		s = data.decode('utf-8', errors='replace')
		if s and s[0] == '\ufeff':
			s = s[1:]
		return s
	except Exception:
		return ''


def _build_display_map(original_path, romanized_path):
	display = {}
	if not (os.path.exists(original_path) and os.path.exists(romanized_path)):
		return display

	try:
		orig_lines = _read_text_any_encoding(original_path).splitlines()
		roman_lines = _read_text_any_encoding(romanized_path).splitlines()
	except Exception:
		return display

	for o_line, r_line in zip(orig_lines, roman_lines):
		o_tokens = [tok for tok in re.split(r'\s+', o_line.strip()) if tok]
		r_tokens = [tok.upper() for tok in re.split(r'\s+', r_line.strip()) if tok]

		if len(o_tokens) != len(r_tokens):
			continue

		for o_tok, r_tok in zip(o_tokens, r_tokens):
			if r_tok:
				display.setdefault(r_tok, o_tok)

	return display


def prep_mlf(trsfile, mlffile, word_dictionary, surround, between):
	"""
	Prepare an input MLF from a transcript, using only words present in the provided dictionary.
	Optionally surround the sentence with tokens and insert a token between words.
	"""
	# Read in the dictionary to ensure all of the words we put in the MLF file are in the dictionary.
	with open(word_dictionary, 'r') as f:
		dictionary = {}
		for line in f.readlines():
			if line != "\n" and line != "":
				dictionary[line.split()[0]] = True

	# Read transcript with robust encoding handling
	content = _read_text_any_encoding(trsfile)
	lines = content.splitlines()

	words = []

	if surround is not None:
		words += surround.split(',')

	# this pattern matches hyphenated words, such as TWENTY-TWO; however, it doesn't work
	# with longer things like SOMETHING-OR-OTHER
	hyphenPat = re.compile(r'([A-Z]+)-([A-Z]+)')

	i = 0
	while i < len(lines):
		txt = lines[i].replace('\n', '')
		txt = txt.replace('{breath}', '{BR}').replace('&lt;noise&gt;', '{NS}')
		txt = txt.replace('{laugh}', '{LG}').replace('{laughter}', '{LG}')
		txt = txt.replace('{cough}', '{CG}').replace('{lipsmack}', '{LS}')

		for pun in [',', '.', ':', ';', '!', '?', '"', '%', '(', ')', '--', '---']:
			txt = txt.replace(pun, '')

		txt = txt.upper()

		# break up any hyphenated words into two separate words
		txt = re.sub(hyphenPat, r'\1 \2', txt)

		txt = txt.split()

		for wrd in txt:
			if wrd in dictionary:
				words.append(wrd)
				if between is not None:
					words.append(between)
			else:
				print("SKIPPING WORD", wrd)

		i += 1

	# Remove the last 'between' token from the end if it exists
	# (though with between_token=None, this won't execute)
	if between is not None and len(words) > 0 and words[-1] == between:
		words.pop()

	if surround is not None:
		words += surround.split(',')

	writeInputMLF(mlffile, words)


def writeInputMLF(mlffile, words):
	with open(mlffile, 'w') as fw:
		fw.write('#!MLF!#\n')
		fw.write('"*/tmp.lab"\n')
		for wrd in words:
			fw.write(wrd + '\n')
		fw.write('.\n')


def readAlignedMLF(mlffile, SR, wave_start):
	"""
	Read a MLF alignment output file with phone and word alignments and return a list of words.
	Each word is a list containing the word label followed by the phones, each phone is a tuple
	(phone, start_time, end_time) with times in seconds.
	sp phones are extracted from words and treated as separate pause intervals.
	"""
	with open(mlffile, 'r') as f:
		lines = [l.rstrip() for l in f.readlines()]

	if len(lines) < 3:
		raise ValueError("Alignment did not complete succesfully.")

	j = 2
	ret = []
	while lines[j] != '.':
		if len(lines[j].split()) == 5:  # start of a word; have a word label?
			# Make a new word list in ret and put the word label at the beginning
			wrd = lines[j].split()[4]
			ret.append([wrd])

		# Append this phone to the latest word (sub-)list
		ph = lines[j].split()[2]
		if SR == 11025:
			st = (float(lines[j].split()[0]) / 10000000.0 + 0.0125) * (11000.0 / 11025.0)
			en = (float(lines[j].split()[1]) / 10000000.0 + 0.0125) * (11000.0 / 11025.0)
		else:
			st = float(lines[j].split()[0]) / 10000000.0 + 0.0125
			en = float(lines[j].split()[1]) / 10000000.0 + 0.0125
		
		# Only add phones with duration > 0
		if st < en:
			ret[-1].append([ph, st + wave_start, en + wave_start])

		j += 1

	# Separate sp from words: if a word ends with sp, move it to a separate entry
	separated_ret = []
	for wrd in ret:
		if len(wrd) > 1 and wrd[-1][0] == 'sp':
			# Word has sp at the end - split it off
			sp_entry = wrd.pop()  # Remove sp from word
			separated_ret.append(wrd)  # Add word without sp
			separated_ret.append(['sp', sp_entry])  # Add sp as separate entry
		else:
			separated_ret.append(wrd)
	
	return separated_ret


def _build_syllable_intervals(word_alignments):
	syllables = []
	for wrd in word_alignments:
		if len(wrd) <= 1:
			continue
		label = wrd[0]
		phones = wrd[1:]
		if label in {'sil', 'sp'}:
			syllables.append([label.upper(), phones[0][1], phones[-1][2]])
			continue

		i = 0
		n = len(phones)
		while i < n:
			syl_phones = []
			syl_start = phones[i][1]

			while i < n and not _is_vowel(phones[i][0]):
				syl_phones.append(phones[i])
				i += 1
				if i < n and _is_vowel(phones[i][0]):
					break
				if i >= n:
					break

			if i < n and _is_vowel(phones[i][0]):
				syl_phones.append(phones[i])
				i += 1
			else:
				if syl_phones:
					syllables.append([
						''.join(p[0] for p in syl_phones if p[0] not in {'sp', 'sil'}).upper() or label.upper(),
						syl_start,
						syl_phones[-1][2]
					])
				break

			consonants = []
			while i < n and not _is_vowel(phones[i][0]):
				consonants.append(phones[i])
				i += 1

			if i < n and consonants:
				carry = consonants[-1:]
				attach = consonants[:-1]
				syl_phones.extend(attach)
				i -= len(carry)
			else:
				syl_phones.extend(consonants)

			syl_end = syl_phones[-1][2]
			syl_label = ''.join(p[0] for p in syl_phones if p[0] not in {'sp', 'sil'}).upper()
			if not syl_label:
				syl_label = label.upper()
			syllables.append([syl_label, syl_start, syl_end])

	return syllables


def _build_utterance_intervals(word_alignments, display_map=None):
	def _display(label):
		if display_map:
			return display_map.get(label, display_map.get(label.upper(), display_map.get(label.lower(), label)))
		return label

	utterances = []
	current_words = []
	current_start = None
	pending_sp_start = None

	for wrd in word_alignments:
		if len(wrd) <= 1:
			continue
		label = wrd[0]
		phones = wrd[1:]

		if label == 'sil':
			sil_end = phones[-1][2]
			if current_words:
				end = pending_sp_start if pending_sp_start is not None else current_words[-1][-1][2]
				start = current_start if current_start is not None else current_words[0][1][1]
				text_tokens = [_display(w[0]) for w in current_words if w[0] not in {'sil', 'sp'}]
				text = ' '.join(t for t in text_tokens if t)
				if text and start < end:
					utterances.append([text, start, end])
				current_words = []
				pending_sp_start = None
			current_start = sil_end
			continue

		if label == 'sp':
			pending_sp_start = phones[0][1]
			continue

		if not current_words and current_start is None:
			current_start = phones[0][1]
		current_words.append(wrd)
		pending_sp_start = None

	# Fallback in case the alignment does not end with sil
	if current_words:
		end = pending_sp_start if pending_sp_start is not None else current_words[-1][-1][2]
		start = current_start if current_start is not None else current_words[0][1][1]
		text_tokens = [_display(w[0]) for w in current_words if w[0] not in {'sil', 'sp'}]
		text = ' '.join(t for t in text_tokens if t)
		if text and start < end:
			utterances.append([text, start, end])

	return utterances


def writeTextGrid(outfile, word_alignments, display_map=None):
	def _display(label):
		if display_map:
			return display_map.get(label, display_map.get(label.upper(), display_map.get(label.lower(), label)))
		return label

	# make the list of just phone alignments
	phons = []
	for wrd in word_alignments:
		phons.extend(wrd[1:])  # skip the word label

	# make the list of just word alignments
	# elements of the form: ["word", ["phone1", st, en], ...]
	wrds = []
	for wrd in word_alignments:
		if len(wrd) == 1:
			continue
		wrds.append([wrd[0], wrd[1][1], wrd[-1][2]])

	# syllable intervals derived from phone-level alignment
	sylls = _build_syllable_intervals(word_alignments)

	# utterance intervals bounded by sil tokens
	utterances = _build_utterance_intervals(word_alignments, display_map)
	tier_start = phons[0][1] if phons else 0.0
	tier_end = phons[-1][-1] if phons else tier_start
	utterance_intervals = []
	cursor = tier_start
	for utt in utterances:
		start, end = utt[1], utt[2]
		if start > cursor:
			utterance_intervals.append(['', cursor, start])
		utterance_intervals.append([utt[0], start, end])
		cursor = end
	if cursor < tier_end:
		utterance_intervals.append(['', cursor, tier_end])
	if not utterance_intervals and tier_start < tier_end:
		utterance_intervals.append(['', tier_start, tier_end])

	# write the phone interval tier
	with open(outfile, 'w') as fw:
		fw.write('File type = "ooTextFile short"\n')
		fw.write('"TextGrid"\n')
		fw.write('\n')
		fw.write(str(phons[0][1]) + '\n')
		fw.write(str(phons[-1][2]) + '\n')
		fw.write('<exists>\n')
		fw.write('4\n')
		fw.write('"IntervalTier"\n')
		fw.write('"phone"\n')
		fw.write(str(phons[0][1]) + '\n')
		fw.write(str(phons[-1][-1]) + '\n')
		fw.write(str(len(phons)) + '\n')
		for k in range(len(phons)):
			fw.write(str(phons[k][1]) + '\n')
			fw.write(str(phons[k][2]) + '\n')
			fw.write('"' + phons[k][0] + '"' + '\n')

		# write the syllable interval tier
		fw.write('"IntervalTier"\n')
		fw.write('"syllable"\n')
		fw.write(str(phons[0][1]) + '\n')
		fw.write(str(phons[-1][-1]) + '\n')
		fw.write(str(len(sylls)) + '\n')
		for syl in sylls:
			fw.write(str(syl[1]) + '\n')
			fw.write(str(syl[2]) + '\n')
			fw.write('"' + syl[0] + '"' + '\n')

		# write the word interval tier
		fw.write('"IntervalTier"\n')
		fw.write('"word"\n')
		fw.write(str(phons[0][1]) + '\n')
		fw.write(str(phons[-1][-1]) + '\n')
		fw.write(str(len(wrds)) + '\n')
		for k in range(len(wrds) - 1):
			fw.write(str(wrds[k][1]) + '\n')
			fw.write(str(wrds[k + 1][1]) + '\n')
			fw.write('"' + _display(wrds[k][0]) + '"' + '\n')

		fw.write(str(wrds[-1][1]) + '\n')
		fw.write(str(phons[-1][2]) + '\n')
		fw.write('"' + _display(wrds[-1][0]) + '"' + '\n')

		# write the utterance interval tier
		fw.write('"IntervalTier"\n')
		fw.write('"utterance"\n')
		fw.write(str(phons[0][1]) + '\n')
		fw.write(str(phons[-1][-1]) + '\n')
		fw.write(str(len(utterance_intervals)) + '\n')
		for utt in utterance_intervals:
			fw.write(str(utt[1]) + '\n')
			fw.write(str(utt[2]) + '\n')
			fw.write('"' + utt[0] + '"' + '\n')


def prep_working_directory():
	os.system("rm -r -f ./tmp")
	os.system("mkdir ./tmp")


def prep_scp(wavfile):
	with open('./tmp/codetr.scp', 'w') as fw:
		fw.write(wavfile + ' ./tmp/tmp.mfc\n')
	with open('./tmp/test.scp', 'w') as fw:
		fw.write('./tmp/tmp.mfc\n')


def create_plp(hcopy_config):
	os.system('HCopy -T 1 -C ' + hcopy_config + ' -S ./tmp/codetr.scp')


def viterbi(input_mlf, word_dictionary, output_mlf, phoneset, hmmdir):
	# MLF includes sil at boundaries, so no -b option needed
	# sp is in dictionary at end of each word
	os.system('HVite -T 1 -a -m -I ' + input_mlf + ' -H ' + hmmdir + '/macros -H ' + hmmdir + '/hmmdefs  -S ./tmp/test.scp -i ' + output_mlf + ' -p 0.0 -s 5.0 ' + word_dictionary + ' ' + phoneset + ' > ./tmp/aligned.results')


def getopt2(name, opts, default=None):
	value = [v for n, v in opts if n == name]
	if len(value) == 0:
		return default
	return value[0]


if __name__ == '__main__':
	try:
		opts, args = getopt.getopt(sys.argv[1:], "r:s:e:", ["model="])

		# get the three mandatory arguments
		if len(args) != 3:
			raise ValueError("Specify wavefile, a transcript file, and an output file!")

		wavfile, trsfile, outfile = args

		sr_override = getopt2("-r", opts, None)
		wave_start = getopt2("-s", opts, "0.0")
		wave_end = getopt2("-e", opts, None)
		# Dictionary includes 'sp' at end of each word
		# Do NOT insert sp in MLF to avoid tee-model conflicts
		surround_token = getopt2("-p", opts, 'sil')
		between_token = getopt2("-b", opts, None)
		display_map = None

		if sr_override is not None:
			try:
				sr_override = int(sr_override)
			except ValueError:
				raise ValueError("-r must be an integer: one of 8000, 11025, or 16000")

		if surround_token is not None and surround_token.strip() == "":
			surround_token = None
		if between_token is not None and between_token.strip() == "":
			between_token = None

		mypath = getopt2("--model", opts, None)
	except Exception:
		print(__doc__)
		(_type, value, _traceback) = sys.exc_info()
		print(value)
		sys.exit(0)

	# If no model directory was said explicitly, get directory containing this script.
	hmmsubdir = ""
	sr_models = None
	if mypath is None:
		mypath = os.path.dirname(os.path.abspath(sys.argv[0])) + "/model"
		hmmsubdir = "FROM-SR"
		# sample rates for which there are acoustic models set up, otherwise
		# the signal must be resampled to one of these rates.
		sr_models = [8000, 11025, 16000]

	if sr_override is not None and sr_models is not None and sr_override not in sr_models:
		raise ValueError("invalid sample rate: not an acoustic model available")

	word_dictionary = "./tmp/dict"
	input_mlf = './tmp/tmp.mlf'
	output_mlf = './tmp/aligned.mlf'

	# create working directory
	prep_working_directory()

	# Helper: detect if transcript contains Hangul (try multiple encodings)
	def contains_hangul(path: str) -> bool:
		try:
			text = _read_text_any_encoding(path)
			for ch in text:
				code = ord(ch)
				if 0xAC00 <= code <= 0xD7A3:
					return True
			return False
		except Exception:
			return False

	# If transcript is in Hangul, convert it to romanized tokens and augment dictionary
	trsfile_for_mlf = trsfile
	if contains_hangul(trsfile):
		print("Detected Hangul transcript; converting and augmenting dictionary...")
		# Prepare temp input in ./tmp
		os.system('cp -f "' + trsfile + '" ./tmp/hangul.txt')
		# 1) Convert sentences to romanized word tokens into ./tmp/sentence_unicode.txt
		#    Note: convert_sentences_unicode.py expects UTF-8 input; our hangul.txt is copied as-is.
		os.system('cd ./tmp && python3 ../bin/convert_sentences_unicode.py hangul.txt')
		trsfile_for_mlf = './tmp/sentence_unicode.txt'
		# 2) Build kdict0/kdict1 in ./tmp from the same Hangul input
		os.system('cd ./tmp && python3 ../bin/han2uniconversion.py hangul.txt')
		os.system('cd ./tmp && python3 ../bin/make_kdict.py')
		# 3) Merge with model dict using add_dict.py inside bin (expects kdict1 in CWD)
		os.system('cp -f ./tmp/kdict1.txt ./bin/kdict1.txt')
		os.system('cd ./bin && python3 add_dict.py')
		# 4) Use the merged dict from bin for this run
		os.system('cp -f ./bin/dict ' + word_dictionary)
		display_map = _build_display_map('./tmp/hangul.txt', trsfile_for_mlf)
	else:
		# Default: start from model dict (+ optional local)
		if os.path.exists("dict.local"):
			os.system("cat " + mypath + "/dict dict.local > " + word_dictionary)
		else:
			os.system("cat " + mypath + "/dict > " + word_dictionary)
		display_map = {}

	# prepare wavefile: do a resampling if necessary
	tmpwav = "./tmp/sound.wav"
	SR = prep_wav(wavfile, tmpwav, sr_override, wave_start, wave_end)

	if hmmsubdir == "FROM-SR":
		hmmsubdir = "/" + str(SR)

	# prepare mlfile (use converted transcript if applicable)
	prep_mlf(trsfile_for_mlf, input_mlf, word_dictionary, surround_token, between_token)

	# prepare scp files
	prep_scp(tmpwav)

	# generate the plp file using a given configuration file for HCopy
	create_plp(mypath + hmmsubdir + '/config')

	# run Viterbi decoding
	print("Running HVite...")
	mpfile = mypath + '/monophones'
	if not os.path.exists(mpfile):
		mpfile = mypath + '/hmmnames'
	viterbi(input_mlf, word_dictionary, output_mlf, mpfile, mypath + hmmsubdir)

	# output the alignment as a Praat TextGrid
	alignments = readAlignedMLF(output_mlf, SR, float(wave_start))
	if display_map is None:
		display_map = {}
	display_map.setdefault('SIL', 'sil')
	display_map.setdefault('SP', 'sp')
	writeTextGrid(outfile, alignments, display_map)