asr.py

import os, sys

from threading import Thread
import gi
gi.require_version('Gst', '1.0')
from gi.repository import GObject, Gst, Gtk, Gdk
GObject.threads_init()
Gdk.threads_init()
Gst.init(None)

import osc

class ASR(Gtk.HBox):
    def __init__(self, asr_config, adaptation=None, partial_hyp_callback = None, hyp_callback = None):
        super(ASR, self).__init__()

        self.text = Gtk.TextView()
        self.text.set_editable(False)
        self.text.set_cursor_visible(False)
        self.buffer = self.text.get_buffer()
        self.text.set_wrap_mode(Gtk.WrapMode.WORD)

        self.scrolled = Gtk.ScrolledWindow()
        self.scrolled.set_policy(Gtk.PolicyType.NEVER, Gtk.PolicyType.ALWAYS)
        self.scrolled.add_with_viewport(self.text)
        self.scrolled.set_size_request(-1, 100)
        self.pack_start(self.scrolled, True, True, 5)

        vbox = Gtk.VBox()
        self.pack_start(vbox, False, False, 5)

        self.button = Gtk.Button("Record")
        self.button.set_size_request(100, -1)
        self.button.set_sensitive(False)
        vbox.pack_start(self.button, True, True, 5)

        #save_state = Gtk.Button("Save state")
        #save_state.connect('clicked', self.save_state)
        #vbox.pack_start(save_state, False, False, 5)
        #load_state = Gtk.Button("Load state")
        #load_state.connect('clicked', self.load_state)
        #vbox.pack_start(load_state, False, False, 5)

        self.button.connect('clicked', self.button_clicked)
        self.text.connect("size-allocate", self.autoscroll)

        self.hyp = []
        self.hyp_callback = hyp_callback
        self.partial_hyp_callback = partial_hyp_callback
        Thread(target=self.init_gst, args=(asr_config, adaptation)).start()

    def init_gst(self, asr_config_file, adaptation):
        """Initialize the speech components"""
        GObject.idle_add(self._started_loading_asr)

        self.pulsesrc = Gst.ElementFactory.make("pulsesrc", "pulsesrc")
        if self.pulsesrc == None:
            print >> sys.stderr, "Error loading pulsesrc GST plugin. You probably need the gstreamer1.0-pulseaudio package"
            sys.exit()	
        self.audioconvert = Gst.ElementFactory.make("audioconvert", "audioconvert")
        self.audioresample = Gst.ElementFactory.make("audioresample", "audioresample")    
        self.asr = Gst.ElementFactory.make("kaldinnet2onlinedecoder", "asr")
        self.fakesink = Gst.ElementFactory.make("fakesink", "fakesink")
        
        if self.asr:
            import config
            config_dict = config.read(asr_config_file)
            for name, value in config_dict.items():
                if name != 'dir' and name != 'name' and not name.startswith('slu') and not name.startswith('xml'):
                    self.asr.set_property(name, value)
        else:
            print >> sys.stderr, "Couldn't create the kaldinnet2onlinedecoder element. "
            if os.environ.has_key("GST_PLUGIN_PATH"):
                print >> sys.stderr, "Have you compiled the Kaldi GStreamer plugin?"
            else:
                print >> sys.stderr, "You probably need to set the GST_PLUGIN_PATH envoronment variable"
                print >> sys.stderr, "Try running: GST_PLUGIN_PATH=../src %s" % sys.argv[0]
                sys.exit();
        
        # initially silence the decoder
        self.asr.set_property("silent", True)
        
        self.pipeline = Gst.Pipeline()
        for element in [self.pulsesrc, self.audioconvert, self.audioresample, self.asr, self.fakesink]:
            self.pipeline.add(element)         
        self.pulsesrc.link(self.audioconvert)
        self.audioconvert.link(self.audioresample)
        self.audioresample.link(self.asr)
        self.asr.link(self.fakesink)    
  
        self.asr.connect('partial-result', self._on_partial_result)
        self.asr.connect('final-result', self._on_final_result)
        self.pipeline.set_state(Gst.State.PLAYING)

        if adaptation != None and adaptation != 'none':
            with open('%s/adaptation-%s.txt' % (config_dict['dir'], adaptation)) as fp:
                data = fp.read()
                #print data
                self.asr.set_property('adaptation-state', data)

        GObject.idle_add(self._finished_loading_asr)

    def load_state(self, *args):
        try:
            with open('state.txt') as fp:
                self.asr.set_property('adaptation-state', fp.read())
                print "ASR: State loaded"
        except:
            print >> sys.stderr, 'failed to load asr state'

    def save_state(self, *args):
        if hasattr(self, 'asr'):
            state = self.asr.get_property('adaptation-state')
            try:
                with open('state.txt', 'w') as fp:
                    fp.write(state)
                print "ASR: State saved"
            except:
                print >> sys.stderr, 'failed to save asr state'

    def _started_loading_asr(self):
        self.button.set_sensitive(False)
        self.button.set_label("Loading")

    def _finished_loading_asr(self):
        self.button.set_label("Record")
        self.button.set_sensitive(True)
        self.button.grab_focus()

    def _on_partial_result(self, asr, hyp):
        hyp = hyp.strip()
        if hyp == '':
            return
        Gdk.threads_enter()
        if len(self.hyp) == 0:
            self.hyp = ['']

        if hyp != self.hyp[-1]:
            osc.client.send_words(len(self.hyp), hyp)

            self.hyp[-1] = hyp
            #print 'PARTIAL', self.hyp
            if self.partial_hyp_callback:
                self.partial_hyp_callback(self.hyp)

            hyp += '...'
            self.insert = self.buffer.get_iter_at_line(self.buffer.get_line_count() - 1)
            self.buffer.delete(self.insert, self.buffer.get_end_iter())
            self.buffer.insert(self.insert, hyp)

        Gdk.threads_leave()
                
    def _on_final_result(self, asr, hyp):
        if len(hyp) > 0 and hyp[-1] == '':
            return
        #print 'FINAL', self.hyp, hyp
        Gdk.threads_enter()
        if len(self.hyp) == 0:
            self.hyp = ['']
        #self.hyp[-1] = hyp
        if self.hyp_callback:
            self.hyp_callback(self.hyp)

        self.insert = self.buffer.get_iter_at_line(self.buffer.get_line_count() - 1)
        self.buffer.delete(self.insert, self.buffer.get_end_iter())
        self.buffer.insert(self.insert, self.hyp[-1] + '\n...')

        self.hyp.append('')
        Gdk.threads_leave()

    def autoscroll(self, *args):
        adj = self.scrolled.get_vadjustment()
        adj.set_value(adj.get_upper() - adj.get_page_size())

    def button_clicked(self, button):
        """Handle button presses."""
        if self.button.get_label() == "Record":
            self.button.set_label("Stop")
            self.asr.set_property("silent", False)
            self.hyp = []
            self.buffer.set_text('...')
        else:
            self.button.set_label("Record")
            self.asr.set_property("silent", True)