astpicotts/picotts.py

#!/usr/bin/env python2

import sys
import os
import shutil
import re
from asterisk import agi
import subprocess
import logging
import hashlib
import uuid

from helpers import get_path, get_sox_ver, detect_format, check_simulate_mode

class PicoTTS:

	cache = False
	cachedir = None
	cache_prefix = 'picotts_'

	tmp_filepaths = []

	def __init__(self, lang=u'fr-FR', tmpdir='/tmp', pico2wave_path=None, sox_path=None, samplerate=None, speed=1, cachedir=None, asterisk_agi=None):
		if not pico2wave_path:
			pico2wave_path=get_path('pico2wave')
			if not pico2wave_path:
				raise Exception('pico2wave not found')
		self.pico2wave_path = pico2wave_path
		if not sox_path:
			sox_path=get_path('sox')
			if not sox_path:
				raise Exception('sox not found')
		self.sox_path = sox_path

		self.asterisk_agi = asterisk_agi
		if not samplerate:
			if not asterisk_agi and not check_simulate_mode():
				logging.warning('You must provide samplerate or asterisk_agi parameter to correctly handle sample rate.')
				(self.samplerate, self.fext) = (8000, 'sln')
			else:
				(self.fext, self.samplerate) = detect_format(asterisk_agi)
		else:
			self.samplerate = samplerate
			self.fext = self.samplerate2fext(samplerate)

		self.speed = speed

		if cachedir:
			self.cachedir = cachedir
			self.check_or_create_cachedir()


		self.tmpdir = tmpdir
		self.lang = lang
		

	def samplerate2fext(self, samplerate):
		# Check/detect sample-rate and final output format
		if samplerate == 8000:
			return "sln"
		elif samplerate == 12000:
			return "sln12"
		elif samplerate == 16000:
			return "sln16"
		elif samplerate == 32000:
			return "sln32"
		elif samplerate == 44100:
			return "sln44"
		elif samplerate == 48000:
			return "sln48"
		else:
			raise Exception('Invalid sample rate value')


	def check_or_create_cachedir(self):
		self.cache = True
		if not os.path.exists(self.cachedir):
			try:
				logging.info('Create cache directory %s' % self.cachedir)
				os.mkdir(self.cachedir)
			except Exception, e:
				logging.warning("Fail to create cache directory (%s) : %s" % (self.cachedir, e))
				logging.info('Disable cache')
				self.cache = False
		else:
			if not os.path.isdir(self.cachedir):
				logging.warning("Cache directory %s is not a directory : disable cache" % self.cachedir)
				self.cache = False
			elif not os.access(self.cachedir, os.W_OK):
				logging.warning("Cache directory %s is not writable : disable cache" % self.cachedir)
				self.cache = False
			else:
				logging.debug("Cache directory %s already exists" % self.cachedir)
		if self.cache:
			logging.debug('Cache is enabled')
		else:
			logging.debug('Cache is disabled')

		return self.cache

	def _getCachePath(self, text, lang=None):
		md5sum=hashlib.md5()
		if isinstance(text, str):
			text = text.decode('utf-8', 'ignore')
		cache_key = u'%s--%s--%s' % (text, (lang or self.lang), self.speed)
		logging.debug('Cache key : "%s"' % cache_key)
		md5sum.update(cache_key.encode('utf-8'))
		cache_md5key = md5sum.hexdigest()
		logging.debug('Cache MD5 key : "%s"' % cache_md5key)
		cache_filename=self.cache_prefix + cache_md5key
		logging.debug('Cache filename : %s' % cache_filename)
		cache_filepath = os.path.join(self.cachedir, cache_filename)
		logging.debug('Cache filepath : %s' % cache_filepath)
		return cache_filepath

	def _getAudioFileFromCache(self, text, lang=None):
		if not self.cache:
			return False
		cache_filepath = self._getCachePath(text, lang=lang) + "." + self.fext
		if os.path.isfile(cache_filepath):
			logging.debug('File already exists in cache. Use it')
			return cache_filepath
		logging.debug('File does not exists in cache.')
		return False

	def getAudioFile(self, text, lang=None):
		if self.cache:
			cache_filepath = self._getAudioFileFromCache(text, lang=lang)
			if cache_filepath:
				return cache_filepath
		
		# Create temp files
		logging.debug('Temporary directory : %s' % self.tmpdir)
		tmpfile = os.path.join(self.tmpdir, 'picotts_%s' % str(uuid.uuid4()))
		tmpwavfile = tmpfile + ".wav"
		logging.debug('Temporary wav file : %s' % tmpwavfile)
		tmpoutfile = tmpfile + "." + self.fext
		logging.debug('Temporary out file : %s' % tmpoutfile)
		
		# Convert text to autio wav file using pico2wave
		cmd = [ self.pico2wave_path, '-l', (lang or self.lang), '-w', tmpwavfile, text ]
		
		try:
			logging.debug('Run command : %s' % cmd)
			result = subprocess.check_output(cmd)
			logging.debug('Command return : %s' % result)
		except subprocess.CalledProcessError, e:
			raise Exception('Fail to convert text to audio file using pico2wave : %s' % e)
		
		# Convert wav file to final output format using sox
		cmd = [ self.sox_path, tmpwavfile, "-q", "-r", str(self.samplerate), "-t", "raw", tmpoutfile ]
		
		# Handle speed change
		if self.speed != 1:
			if get_sox_ver(sox_path=self.sox_path) >= 14:
				cmd += ['tempo', '-s', str(self.speed)]
			else:
				cmd += ['stretch', str(1/self.speed), "80"]
		
		try:
			logging.debug('Run command : %s' % cmd)
			result = subprocess.check_output(cmd)
			logging.debug('Command return : %s' % result)
		except subprocess.CalledProcessError, e:
			logging.fatal('Fail to convert text to audio file using pico2wave : %s' % e)
			sys.exit(1)
		
		try:
			logging.debug('Remove tmp wav file')
			os.remove(tmpwavfile)
		except Exception, e:
			logging.warning('Fail to remove temporary WAV audio file')

		# Move audio file in cache
		if self.cache:
			cache_filepath = self._getCachePath(text, lang=lang) + "." + self.fext
			try:
				logging.debug('Move audio file in cache directory')
				shutil.move(tmpoutfile, cache_filepath)
			except Exception,e:
				logging.warning('Fail to move audio file in cache directory : %s' % e)
			return cache_filepath
		else:
			self.tmp_filepaths.append(tmpoutfile)
			logging.debug('Cache disabled. Directly play tmp file.')
		return tmpoutfile

	def clean_tmp(self):
		try:
			logging.debug('Clean temporaries files and directory')
			for filepath in self.tmp_filepaths:
				if os.path.exists(filepath):
					os.remove(filepath)
		except Exception as e:
			logging.warning('Fail to remove temporaries files and directory : %s' % e)

if __name__ == "__main__":
	from optparse import OptionParser
	from helpers import playback, enable_simulate_mode
	import traceback

	default_logfile = '/var/log/asterisk/picotts.log'
	default_lang = 'fr-FR'
	default_intkey = ""
	default_result_varname = "USER_INPUT"
	default_speed = 1
	default_cachedir = '/tmp/'
	default_read_timeout = 3000
	default_read_maxdigits = 20
	
	any_intkeys = "0123456789#*"
	
	#######
	# RUN #
	#######
	
	# Options parser
	parser = OptionParser()
	
	parser.add_option('-d', '--debug',
			  action="store_true",
			  dest="debug",
			  help="Enable debug mode")
	
	parser.add_option('-v', '--verbose',
			  action="store_true",
			  dest="verbose",
			  help="Enable verbose mode")
	
	parser.add_option('--simulate',
			  action="store_true",
			  dest="simulate",
			  help="Simulate AGI mode")
	
	parser.add_option('--simulate-play',
			  action="store_true",
			  dest="simulate_play",
			  help="Simulate mode : play file using mplayer")
	
	parser.add_option('-r', '--read',
			  action="store_true",
			  dest="read",
			  help="Enable read mode")
	
	parser.add_option('-t', '--read-timeout',
			  action="store", type="int",
			  dest="read_timeout", default=default_read_timeout,
			  help="Read timeout in ms (Default : %i)" % default_read_timeout)
	
	parser.add_option('-m', '--read-max-digits',
			  action="store", type="int",
			  dest="read_maxdigits", default=default_read_maxdigits,
			  help="Read max digits (Default : %i)" % default_read_maxdigits)
	
	parser.add_option('-n', '--name',
			  action="store", type="string",
			  dest="varname", default=default_result_varname,
			  help="User input result variable name (Default : %s)" % default_result_varname)
	
	parser.add_option('-L', '--log-file',
			  action="store", type="string",
			  dest="logfile", default=default_logfile,
			  help="pico2wave path (Default : %s)" % default_logfile)
	
	parser.add_option('-l', '--lang',
			  action="store", type="string",
			  dest="lang", default=default_lang,
			  help="Language (Default : %s)" % default_lang)
	
	parser.add_option('-i', '--intkey',
			  action="store", type="string",
			  dest="intkey", default=default_intkey,
			  help="Interrupt key(s) (Default : No)")
	
	parser.add_option('-s', '--speed',
			  action="store", type="float",
			  dest="speed", default=default_speed,
			  help="Speed factor (Default : %i)" % default_speed)
	
	parser.add_option('-S', '--sample-rate',
			  action="store", type="int",
			  dest="samplerate",
			  help="Sample rate (Default : auto-detect)")
	
	parser.add_option('-c', '--cache',
			  action="store_true",
			  dest="cache",
			  help="Enable cache")
	
	parser.add_option('-C', '--cache-dir',
			  action="store", type="string",
			  dest="cachedir", default=default_cachedir,
			  help="Cache directory path (Default : %s)" % default_cachedir)
	
	parser.add_option('--sox-path',
			  action="store", type="string",
			  dest="sox_path",
			  help="sox path (Default : auto-detec in PATH)")
	
	parser.add_option('--pico2wave-path',
			  action="store", type="string",
			  dest="pico2wave_path",
			  help="pico2wave path (Default : auto-detec in PATH)")
	
	(options, args) = parser.parse_args()

	try:	
		# Enable logs
		logformat = '%(levelname)s - %(message)s'
		if options.simulate:
			logging.basicConfig(format=logformat, level=logging.DEBUG)
			enable_simulate_mode()
		else:
			if options.debug:
				loglevel = logging.DEBUG
			elif options.verbose:
				loglevel = logging.INFO
			else:
				loglevel = logging.WARNING
			logging.basicConfig(filename=options.logfile, level=loglevel,
				    format=logformat)
		
		text=" ".join(args).strip()
		if len(text) == 0:
			usage(msg="You provide text to say as first parameter.")
		
		# Valid intkey parameter
		if options.intkey == "any":
			options.intkey = any_intkeys
		elif not re.match('^[0-9#*]*$', options.intkey):
			logging.warning('Invalid interrupt key(s) provided ("%s"), use any.' % options.intkey)
			options.intkey = any_intkeys
		
		if options.speed <= 0:
			logging.warning('Invalid speed provided, use default')
			options.speed=default_speed
		
		logging.debug('Call parameters (lang = {lang}, intkey = {intkey}, speed = {speed} and text :\n{text}'.format(
			lang=options.lang, intkey=options.intkey,
			speed=options.speed, text=text)
		)
		
		picotts_args = {}
		# Start Asterisk AGI client
		if not options.simulate:
			a = agi.AGI()
			picotts_args['asterisk_agi'] = a
		else:
			a = None
		
		if options.cache:
			picotts_args['cachedir']=options.cachedir
		
		picotts = PicoTTS(lang = options.lang, speed=options.speed, **picotts_args)
		
		filepath = picotts.getAudioFile(text)
		
		# Playback message
		result = playback(a, filepath, simulate_play=options.simulate_play,
		         intkey=options.intkey, read=options.read,
		         read_timeout=options.read_timeout,
		         read_maxdigits=options.read_maxdigits)
		
		logging.debug('User enter : "%s"' % result)
		if options.simulate:
			logging.debug('Simulate mode : set variable %s to "%s"' % (options.varname, result))
		else:
			logging.info('Set variable %s to "%s"' % (options.varname, result))
			a.set_variable(options.varname, result)

		picotts.clean_tmp()		
	except agi.AGIAppError as e:
		logging.info('An AGI error stop script : %s' % e)
		if 'picotts' in globals():
			picotts.clean_tmp()
		sys.exit(0)
	except Exception as e:
		logging.error(traceback.format_exc())
		if 'picotts' in globals():
			picotts.clean_tmp()
		sys.exit(1)
	sys.exit(0)