Merge pull request #176 from Holzhaus/improved-stt-module

Improved stt module

2 files changed, 156 insertions, 90 deletions

diff --git a/client/mic.py b/client/mic.py
index 0351592..2e8d603 100644
--- a/client/mic.py
+++ b/client/mic.py
@@ -4,11 +4,13 @@
 """
 
 import os
-from wave import open as open_audio
+import tempfile
+import wave
 import audioop
 import pyaudio
 import alteration
 import jasperpath
+from stt import TranscriptionMode
 
 class Mic:
 
@@ -80,7 +82,6 @@ class Mic:
         """
 
         THRESHOLD_MULTIPLIER = 1.8
-        AUDIO_FILE = "passive.wav"
         RATE = 16000
         CHUNK = 1024
 
@@ -154,15 +155,17 @@ class Mic:
         stream.stop_stream()
         stream.close()
         audio.terminate()
-        write_frames = open_audio(AUDIO_FILE, 'wb')
-        write_frames.setnchannels(1)
-        write_frames.setsampwidth(audio.get_sample_size(pyaudio.paInt16))
-        write_frames.setframerate(RATE)
-        write_frames.writeframes(''.join(frames))
-        write_frames.close()
-
-        # check if PERSONA was said
-        transcribed = self.passive_stt_engine.transcribe(AUDIO_FILE, PERSONA_ONLY=True)
+        
+        with tempfile.NamedTemporaryFile(mode='w+b') as f:
+            wav_fp = wave.open(f, 'wb')
+            wav_fp.setnchannels(1)
+            wav_fp.setsampwidth(audio.get_sample_size(pyaudio.paInt16))
+            wav_fp.setframerate(RATE)
+            wav_fp.writeframes(''.join(frames))
+            wav_fp.close()
+            f.seek(0)
+            # check if PERSONA was said
+            transcribed = self.passive_stt_engine.transcribe(f, mode=TranscriptionMode.KEYWORD)
 
         if PERSONA in transcribed:
             return (THRESHOLD, PERSONA)
@@ -187,18 +190,10 @@ class Mic:
             Returns a list of the matching options or None
         """
 
-        AUDIO_FILE = "active.wav"
         RATE = 16000
         CHUNK = 1024
         LISTEN_TIME = 12
 
-        # user can request pre-recorded sound
-        if not LISTEN:
-            if not os.path.exists(AUDIO_FILE):
-                return None
-
-            return self.active_stt_engine.transcribe(AUDIO_FILE)
-
         # check if no threshold provided
         if THRESHOLD == None:
             THRESHOLD = self.fetchThreshold()
@@ -238,14 +233,18 @@ class Mic:
         stream.stop_stream()
         stream.close()
         audio.terminate()
-        write_frames = open_audio(AUDIO_FILE, 'wb')
-        write_frames.setnchannels(1)
-        write_frames.setsampwidth(audio.get_sample_size(pyaudio.paInt16))
-        write_frames.setframerate(RATE)
-        write_frames.writeframes(''.join(frames))
-        write_frames.close()
 
-        return self.active_stt_engine.transcribe(AUDIO_FILE, MUSIC=MUSIC)
+        with tempfile.SpooledTemporaryFile(mode='w+b') as f:
+            wav_fp = wave.open(f, 'wb')
+            wav_fp.setnchannels(1)
+            wav_fp.setsampwidth(audio.get_sample_size(pyaudio.paInt16))
+            wav_fp.setframerate(RATE)
+            wav_fp.writeframes(''.join(frames))
+            wav_fp.close()
+            f.seek(0)
+            mode = TranscriptionMode.MUSIC if MUSIC else TranscriptionMode.NORMAL
+            transcribed = self.active_stt_engine.transcribe(f, mode=mode)
+        return transcribed
 
     def say(self, phrase, OPTIONS=" -vdefault+m3 -p 40 -s 160 --stdout > say.wav"):
         # alter phrase before speaking
diff --git a/client/stt.py b/client/stt.py
index fc1bd1f..899775b 100644
--- a/client/stt.py
+++ b/client/stt.py
@@ -2,9 +2,11 @@
 # -*- coding: utf-8-*-
 import os
 import traceback
+import wave
 import json
 import tempfile
 import logging
+from abc import ABCMeta, abstractmethod
 import requests
 import yaml
 
@@ -12,12 +14,37 @@ import yaml
 The default Speech-to-Text implementation which relies on PocketSphinx.
 """
 
+class TranscriptionMode:
+    NORMAL, KEYWORD, MUSIC = range(3)
 
-class PocketSphinxSTT(object):
+class AbstractSTTEngine(object):
+    """
+    Generic parent class for all STT engines
+    """
+
+    __metaclass__ = ABCMeta
+    
+    @classmethod
+    def get_config(cls):
+        return {}
+
+    @classmethod
+    @abstractmethod
+    def is_available(cls):
+        return True
+
+    @abstractmethod
+    def transcribe(self, fp, mode=TranscriptionMode.NORMAL):
+        pass
+
+class PocketSphinxSTT(AbstractSTTEngine):
+
+    SLUG = 'sphinx'
 
     def __init__(self, lmd="languagemodel.lm", dictd="dictionary.dic",
                  lmd_persona="languagemodel_persona.lm", dictd_persona="dictionary_persona.dic",
-                 lmd_music=None, dictd_music=None, **kwargs):
+                 lmd_music=None, dictd_music=None,
+                 hmm_dir="/usr/local/share/pocketsphinx/model/hmm/en_US/hub4wsj_sc_8k"):
         """
         Initiates the pocketsphinx instance.
 
@@ -37,38 +64,51 @@ class PocketSphinxSTT(object):
         except:
             import pocketsphinx as ps
 
-        hmm_dir = None
+        self._logfiles = {}
+        with tempfile.NamedTemporaryFile(prefix='psdecoder_music_', suffix='.log', delete=False) as f:
+            self._logfiles[TranscriptionMode.MUSIC] = f.name
+        with tempfile.NamedTemporaryFile(prefix='psdecoder_keyword_', suffix='.log', delete=False) as f:
+            self._logfiles[TranscriptionMode.KEYWORD] = f.name
+        with tempfile.NamedTemporaryFile(prefix='psdecoder_normal_', suffix='.log', delete=False) as f:
+            self._logfiles[TranscriptionMode.NORMAL] = f.name
+
+        self._decoders = {}
+        if lmd_music and dictd_music:
+            self._decoders[TranscriptionMode.MUSIC] = ps.Decoder(hmm=hmm_dir, lm=lmd_music, dict=dictd_music, logfn=self._logfiles[TranscriptionMode.MUSIC])
+        self._decoders[TranscriptionMode.KEYWORD]  = ps.Decoder(hmm=hmm_dir, lm=lmd_persona, dict=dictd_persona, logfn=self._logfiles[TranscriptionMode.KEYWORD])
+        self._decoders[TranscriptionMode.NORMAL] = ps.Decoder(hmm=hmm_dir, lm=lmd, dict=dictd, logfn=self._logfiles[TranscriptionMode.NORMAL])
 
+    def __del__(self):
+        for filename in self._logfiles.values():
+            os.remove(filename)
+
+    @classmethod
+    def get_config(cls): #FIXME: Replace this as soon as we have a config module
+        config = {}
+        # HMM dir
         # Try to get hmm_dir from config
         profile_path = os.path.join(os.path.dirname(__file__), 'profile.yml')
         if os.path.exists(profile_path):
             with open(profile_path, 'r') as f:
                 profile = yaml.safe_load(f)
-                if 'pocketsphinx' in profile and 'hmm_dir' in profile['pocketsphinx']:
-                    hmm_dir = profile['pocketsphinx']['hmm_dir']
-
-        if not hmm_dir:
-            hmm_dir = "/usr/local/share/pocketsphinx/model/hmm/en_US/hub4wsj_sc_8k"
-
-        with tempfile.NamedTemporaryFile(prefix='psdecoder_music_', suffix='.log', delete=False) as f:
-            self.logfile_music = f.name
-        with tempfile.NamedTemporaryFile(prefix='psdecoder_persona_', suffix='.log', delete=False) as f:
-            self.logfile_persona = f.name
-        with tempfile.NamedTemporaryFile(prefix='psdecoder_default_', suffix='.log', delete=False) as f:
-            self.logfile_default = f.name
-
-        if lmd_music and dictd_music:
-            self.speechRec_music = ps.Decoder(hmm=hmm_dir, lm=lmd_music, dict=dictd_music, logfn=self.logfile_music)
-        self.speechRec_persona = ps.Decoder(
-            hmm=hmm_dir, lm=lmd_persona, dict=dictd_persona, logfn=self.logfile_persona)
-        self.speechRec = ps.Decoder(hmm=hmm_dir, lm=lmd, dict=dictd, logfn=self.logfile_default)
-
-    def __del__(self):
-        os.remove(self.logfile_music)
-        os.remove(self.logfile_persona)
-        os.remove(self.logfile_default)
+                if 'pocketsphinx' in profile:
+                    if 'hmm_dir' in profile['pocketsphinx']:
+                        config['hmm_dir'] = profile['pocketsphinx']['hmm_dir']
+                    if 'lmd' in profile['pocketsphinx']:
+                        config['lmd'] = profile['pocketsphinx']['lmd']
+                    if 'dictd' in profile['pocketsphinx']:
+                        config['dictd'] = profile['pocketsphinx']['dictd']
+                    if 'lmd_persona' in profile['pocketsphinx']:
+                        config['lmd_persona'] = profile['pocketsphinx']['lmd_persona']
+                    if 'dictd_persona' in profile['pocketsphinx']:
+                        config['dictd_persona'] = profile['pocketsphinx']['dictd_persona']
+                    if 'lmd_music' in profile['pocketsphinx']:
+                        config['lmd'] = profile['pocketsphinx']['lmd_music']
+                    if 'dictd_music' in profile['pocketsphinx']:
+                        config['dictd_music'] = profile['pocketsphinx']['dictd_music']
+        return config
 
-    def transcribe(self, audio_file_path, PERSONA_ONLY=False, MUSIC=False):
+    def transcribe(self, fp, mode=TranscriptionMode.NORMAL):
         """
         Performs STT, transcribing an audio file and returning the result.
 
@@ -77,30 +117,27 @@ class PocketSphinxSTT(object):
         PERSONA_ONLY -- if True, uses the 'Persona' language model and dictionary
         MUSIC -- if True, uses the 'Music' language model and dictionary
         """
+        decoder = self._decoders[mode]
 
-        wavFile = file(audio_file_path, 'rb')
-        wavFile.seek(44)
+        fp.seek(44)
 
-        if MUSIC:
-            self.speechRec_music.decode_raw(wavFile)
-            result = self.speechRec_music.get_hyp()
-            with open(self.logfile_music, 'r+') as f:
+        # FIXME: Can't use the Decoder.decode_raw() here, because
+        # pocketsphinx segfaults with tempfile.SpooledTemporaryFile()
+        data = fp.read()
+        decoder.start_utt()
+        decoder.process_raw(data, False, True)
+        decoder.end_utt()
+        
+        result = decoder.get_hyp()
+        with open(self._logfiles[mode], 'r+') as f:
+                if mode == TranscriptionMode.KEYWORD:
+                    modename = "[KEYWORD]"
+                elif mode == TranscriptionMode.MUSIC:
+                    modename = "[MUSIC]"
+                else:
+                    modename = "[NORMAL]"
                 for line in f:
-                    self._logger.debug("speechRec_music %s", line.strip())
-                f.truncate()
-        elif PERSONA_ONLY:
-            self.speechRec_persona.decode_raw(wavFile)
-            result = self.speechRec_persona.get_hyp()
-            with open(self.logfile_persona, 'r+') as f:
-                for line in f:
-                    self._logger.debug("speechRec_persona %s", line.strip())
-                f.truncate()
-        else:
-            self.speechRec.decode_raw(wavFile)
-            result = self.speechRec.get_hyp()
-            with open(self.logfile_default, 'r+') as f:
-                for line in f:
-                    self._logger.debug("speechRec_default %s", line.strip())
+                    self._logger.debug("%s %s", modename, line.strip())
                 f.truncate()
 
         print "==================="
@@ -109,6 +146,10 @@ class PocketSphinxSTT(object):
 
         return [result[0]]
 
+    @classmethod
+    def is_available(cls):
+        return (pkgutil.get_loader('pocketsphinx') is not None and os.path.exits(cls._get_hmm_dir()))
+
 """
 Speech-To-Text implementation which relies on the Google Speech API.
 
@@ -134,19 +175,34 @@ Excerpt from sample profile.yml:
 """
 
 
-class GoogleSTT(object):
+class GoogleSTT(AbstractSTTEngine):
 
-    RATE = 16000
+    SLUG = 'google'
 
-    def __init__(self, api_key, **kwargs):
+    def __init__(self, api_key=None): #FIXME: get init args from config
         """
         Arguments:
         api_key - the public api key which allows access to Google APIs
         """
+        if not api_key:
+            raise ValueError("No Google API Key given")
         self.api_key = api_key
         self.http = requests.Session()
 
-    def transcribe(self, audio_file_path, PERSONA_ONLY=False, MUSIC=False):
+    @classmethod
+    def get_config(cls): #FIXME: Replace this as soon as we have a config module
+        config = {}
+        # HMM dir
+        # Try to get hmm_dir from config
+        profile_path = os.path.join(os.path.dirname(__file__), 'profile.yml')
+        if os.path.exists(profile_path):
+            with open(profile_path, 'r') as f:
+                profile = yaml.safe_load(f)
+                if 'keys' in profile and 'GOOGLE_SPEECH' in profile['keys']:
+                    config['api_key'] = profile['keys']['GOOGLE_SPEECH']
+        return config
+
+    def transcribe(self, fp, mode=TranscriptionMode.NORMAL):
         """
         Performs STT via the Google Speech API, transcribing an audio file and returning an English
         string.
@@ -154,15 +210,18 @@ class GoogleSTT(object):
         Arguments:
         audio_file_path -- the path to the .wav file to be transcribed
         """
+
+        wav = wave.open(fp, 'rb')
+        frame_rate = wav.getframerate()
+        wav.close()
+
         url = "https://www.google.com/speech-api/v2/recognize?output=json&client=chromium&key=%s&lang=%s&maxresults=6&pfilter=2" % (
             self.api_key, "en-us")
 
-        wav = open(audio_file_path, 'rb')
-        data = wav.read()
-        wav.close()
+        data = fp.read()
 
         try:
-            headers = {'Content-type': 'audio/l16; rate=%s' % GoogleSTT.RATE}
+            headers = {'Content-type': 'audio/l16; rate=%s' % frame_rate}
             response = self.http.post(url, data=data, headers=headers)
             response.encoding = 'utf-8'
             response_read = response.text
@@ -181,6 +240,10 @@ class GoogleSTT(object):
         except Exception:
             traceback.print_exc()
 
+    @classmethod
+    def is_available(cls):
+        return True
+
 """
 Returns a Speech-To-Text engine.
 
@@ -191,13 +254,17 @@ Arguments:
 engine_type - one of "sphinx" or "google"
 kwargs - keyword arguments passed to the constructor of the STT engine
 """
+def get_engines():
+    return [stt_engine for stt_engine in AbstractSTTEngine.__subclasses__() if hasattr(stt_engine, 'SLUG') and stt_engine.SLUG]
 
-
-def newSTTEngine(engine_type, **kwargs):
-    t = engine_type.lower()
-    if t == "sphinx":
-        return PocketSphinxSTT(**kwargs)
-    elif t == "google":
-        return GoogleSTT(**kwargs)
+def newSTTEngine(stt_engine, **kwargs):
+    selected_engines = filter(lambda engine: hasattr(engine, "SLUG") and engine.SLUG == stt_engine, get_engines())
+    if len(selected_engines) == 0:
+        raise ValueError("No STT engine found for slug '%s'" % stt_engine)
     else:
-        raise ValueError("Unsupported STT engine type: " + engine_type)
+        if len(selected_engines) > 1:
+            print("WARNING: Multiple STT engines found for slug '%s'. This is most certainly a bug." % stt_engine)
+        engine = selected_engines[0]
+        if not engine.is_available():
+            raise ValueError("STT engine '%s' is not available (due to missing dependencies, missing dependencies, etc.)" % stt_engine)
+        return engine(**engine.get_config())