text:
import os
import shutil
import glob
import subprocess
import re
import math
import numpy as np
from tqdm import tqdm

from ..helper import thirdparty_binary, make_path_safe

from ..config import (MonophoneConfig, TriphoneConfig, TriphoneFmllrConfig,
                      LdaMlltConfig, DiagUbmConfig, iVectorExtractorConfig,
                      NnetBasicConfig)

from ..multiprocessing import (align, mono_align_equal, compile_train_graphs,
                               acc_stats, tree_stats, convert_alignments,
                               convert_ali_to_textgrids, calc_fmllr,
                               lda_acc_stats,
                               calc_lda_mllt, gmm_gselect, acc_global_stats,
                               gauss_to_post, acc_ivector_stats, get_egs,
                               get_lda_nnet, nnet_train_trans, nnet_train,
                               nnet_align, nnet_get_align_feats, extract_ivectors,
                               compute_prob, get_average_posteriors, relabel_egs)
from ..accuracy_graph import get_accuracy_graph

from ..exceptions import NoSuccessfulAlignments

from .. import __version__

from ..config import TEMP_DIR

class BaseAligner(object):
    '''
    Base aligner class for common aligner functions

    Parameters
    ----------
    corpus : :class:`~aligner.corpus.Corpus`
        Corpus object for the dataset
    dictionary : :class:`~aligner.dictionary.Dictionary`
        Dictionary object for the pronunciation dictionary
    output_directory : str
        Path to export aligned TextGrids
    temp_directory : str, optional
        Specifies the temporary directory root to save files need for Kaldi.
        If not specified, it will be set to ``~/Documents/MFA``
    num_jobs : int, optional
        Number of processes to use, defaults to 3
    call_back : callable, optional
        Specifies a call back function for alignment
    mono_params : :class:`~aligner.config.MonophoneConfig`, optional
        Monophone training parameters to use, if different from defaults
    tri_params : :class:`~aligner.config.TriphoneConfig`, optional
        Triphone training parameters to use, if different from defaults
    tri_fmllr_params : :class:`~aligner.config.TriphoneFmllrConfig`, optional
        Speaker-adapted triphone training parameters to use, if different from defaults
    '''

    def __init__(self, corpus, dictionary, output_directory,
                 temp_directory=None, num_jobs=3, call_back=None,
                 mono_params=None, tri_params=None,
                 tri_fmllr_params=None, lda_mllt_params=None,
                 diag_ubm_params=None, ivector_extractor_params=None,
                 nnet_basic_params=None,
                 debug=False, nnet=False):
        self.nnet = nnet
        if mono_params is None:
            mono_params = {}
        if tri_params is None:
            tri_params = {}
        if tri_fmllr_params is None:
            tri_fmllr_params = {}

        if lda_mllt_params is None:
            lda_mllt_params = {}
        if diag_ubm_params is None:
            diag_ubm_params = {}
        if ivector_extractor_params is None:
            ivector_extractor_params = {}
        if nnet_basic_params is None:
            nnet_basic_params = {}

        self.mono_config = MonophoneConfig(**mono_params)
        self.tri_config = TriphoneConfig(**tri_params)
        self.tri_fmllr_config = TriphoneFmllrConfig(**tri_fmllr_params)

        self.lda_mllt_config = LdaMlltConfig(**lda_mllt_params)
        self.diag_ubm_config = DiagUbmConfig(**diag_ubm_params)
        self.ivector_extractor_config = iVectorExtractorConfig(**ivector_extractor_params)
        self.nnet_basic_config = NnetBasicConfig(**nnet_basic_params)

        self.corpus = corpus
        self.dictionary = dictionary
        self.output_directory = output_directory
        if self.corpus.num_jobs != num_jobs:
            num_jobs = self.corpus.num_jobs
        self.num_jobs = num_jobs
        if not temp_directory:
            temp_directory = TEMP_DIR
        self.temp_directory = temp_directory
        self.call_back = call_back
        if self.call_back is None:
            self.call_back = print
        self.verbose = False
        self.debug = debug
        self.setup()

    def setup(self):
        self.dictionary.write()
        self.corpus.initialize_corpus(self.dictionary)
        print(self.corpus.speaker_utterance_info())

    @property
    def meta(self):
        data = {'phones':sorted(self.dictionary.nonsil_phones),
                'version': __version__,
                'architecture':'gmm-hmm',
                'features':'mfcc+deltas',
                }
        return data

    @property
    def mono_directory(self):
        return os.path.join(self.temp_directory, 'mono')

    @property
    def mono_final_model_path(self):
        return os.path.join(self.mono_directory, 'final.mdl')

    @property
    def mono_ali_directory(self):
        return os.path.join(self.temp_directory, 'mono_ali')

    @property
    def tri_directory(self):
        return os.path.join(self.temp_directory, 'tri')

    @property
    def tri_ali_directory(self):
        return os.path.join(self.temp_directory, 'tri_ali')

    @property
    def tri_final_model_path(self):
        return os.path.join(self.tri_directory, 'final.mdl')

    @property
    def tri_fmllr_directory(self):
        return os.path.join(self.temp_directory, 'tri_fmllr')

    @property
    def tri_fmllr_ali_directory(self):
        return os.path.join(self.temp_directory, 'tri_fmllr_ali')

    @property
    def tri_fmllr_final_model_path(self):
        return os.path.join(self.tri_fmllr_directory, 'final.mdl')

    # Beginning of nnet properties
    @property
    def lda_mllt_directory(self):
        return os.path.join(self.temp_directory, 'lda_mllt')

    @property
    def lda_mllt_ali_directory(self):
        return os.path.join(self.temp_directory, 'lda_mllt_ali')

    @property
    def lda_mllt_final_model_path(self):
        return os.path.join(self.lda_mllt_directory, 'final.mdl')

    @property
    def diag_ubm_directory(self):
        return os.path.join(self.temp_directory, 'diag_ubm')

    @property
    def diag_ubm_final_model_path(self):
        return os.path.join(self.diag_ubm_directory, 'final.dubm')

    @property
    def ivector_extractor_directory(self):
        return os.path.join(self.temp_directory, 'ivector_extractor')

    @property
    def ivector_extractor_final_model_path(self):
        return os.path.join(self.ivector_extractor_directory, 'final.ie')

    @property
    def extracted_ivector_directory(self):
        return os.path.join(self.temp_directory, 'extracted_ivector')

    @property
    def nnet_basic_directory(self):
        return os.path.join(self.temp_directory, 'nnet_basic')

    @property
    def nnet_basic_ali_directory(self):
        return os.path.join(self.temp_directory, 'nnet_basic_ali')

    @property
    def nnet_basic_final_model_path(self):
        return os.path.join(self.nnet_basic_directory, 'final.mdl')

    # End of nnet properties

    def export_textgrids(self):
        '''
        Export a TextGrid file for every sound file in the dataset
        '''
        if os.path.exists(self.nnet_basic_final_model_path):
            model_directory = self.nnet_basic_directory
        elif os.path.exists(self.tri_fmllr_final_model_path):
            model_directory = self.tri_fmllr_directory
        elif os.path.exists(self.tri_final_model_path):
            model_directory = self.tri_directory
        elif os.path.exists(self.mono_final_model_path):
            model_directory = self.mono_directory
        # Debugging
        #print("EXPORTING TO:", self.output_directory)
        convert_ali_to_textgrids(self.output_directory, model_directory, self.dictionary,
                                 self.corpus, self.num_jobs)

    def get_num_gauss_mono(self):
        '''
        Get the number of gaussians for a monophone model
        '''
        with open(os.devnull, 'w') as devnull:
            proc = subprocess.Popen([thirdparty_binary('gmm-info'),
                                     '--print-args=false',
                                     os.path.join(self.mono_directory, '0.mdl')],
                                    stderr=devnull,
                                    stdout=subprocess.PIPE)
            stdout, stderr = proc.communicate()
            num = stdout.decode('utf8')
            matches = re.search(r'gaussians (\d+)', num)
            num = int(matches.groups()[0])
        return num

    def _align_si(self, fmllr=False, lda_mllt=False, feature_name=None):
        '''
        Generate an alignment of the dataset
        '''
        if fmllr and os.path.exists(self.tri_fmllr_final_model_path):
            model_directory = self.tri_fmllr_directory
            output_directory = self.tri_fmllr_ali_directory
            config = self.tri_fmllr_config

        elif fmllr:     # First pass with fmllr, final path doesn't exist yet
            model_directory = self.tri_directory
            output_directory = self.tri_fmllr_ali_directory
            config = self.tri_fmllr_config

        elif lda_mllt and os.path.exists(self.lda_mllt_final_model_path):
            model_directory = self.lda_mllt_directory
            output_directory = self.lda_mllt_ali_directory
            config = self.lda_mllt_config

        elif lda_mllt:  # First pass with LDA + MLLT, final path doesn't exist yet
            model_directory = self.tri_fmllr_directory
            output_directory = self.lda_mllt_ali_directory
            config = self.lda_mllt_config
        elif os.path.exists(self.tri_final_model_path):
            model_directory = self.tri_directory
            output_directory = self.tri_ali_directory
            config = self.tri_config
        elif os.path.exists(self.mono_final_model_path):
            model_directory = self.mono_directory
            output_directory = self.mono_ali_directory
            config = self.mono_config

        optional_silence = self.dictionary.optional_silence_csl
        oov = self.dictionary.oov_int

        log_dir = os.path.join(output_directory, 'log')
        os.makedirs(self.tri_fmllr_ali_directory, exist_ok=True)
        os.makedirs(self.lda_mllt_ali_directory, exist_ok=True)

        os.makedirs(log_dir, exist_ok=True)

        shutil.copyfile(os.path.join(model_directory, 'tree'),
                        os.path.join(output_directory, 'tree'))
        shutil.copyfile(os.path.join(model_directory, 'final.mdl'),
                        os.path.join(output_directory, '0.mdl'))

        shutil.copyfile(os.path.join(model_directory, 'final.occs'),
                        os.path.join(output_directory, '0.occs'))

        feat_type = 'delta'

        compile_train_graphs(output_directory, self.dictionary.output_directory,
                             self.corpus.split_directory, self.num_jobs, debug=self.debug)

        align(0, output_directory, self.corpus.split_directory,
              optional_silence, self.num_jobs, config, feature_name=feature_name)
        shutil.copyfile(os.path.join(output_directory, '0.mdl'), os.path.join(output_directory, 'final.mdl'))
        shutil.copyfile(os.path.join(output_directory, '0.occs'), os.path.join(output_directory, 'final.occs'))

        # Debugging
        #print("checking:", os.path.exists(os.path.join(self.tri_fmllr_directory, 'final.mdl')))
        if output_directory == self.tri_fmllr_ali_directory:
            os.makedirs(self.tri_fmllr_directory, exist_ok=True)
            shutil.copyfile(os.path.join(output_directory, '0.mdl'), os.path.join(self.tri_fmllr_directory, 'final.mdl'))
            shutil.copyfile(os.path.join(output_directory, '0.occs'), os.path.join(self.tri_fmllr_directory, 'final.occs'))
        elif output_directory == self.lda_mllt_ali_directory:
            os.makedirs(self.lda_mllt_directory, exist_ok=True)
            shutil.copyfile(os.path.join(output_directory, '0.mdl'), os.path.join(self.lda_mllt_directory, 'final.mdl'))
            shutil.copyfile(os.path.join(output_directory, '0.occs'), os.path.join(self.lda_mllt_directory, 'final.occs'))

    def parse_log_directory(self, directory, iteration):
        '''
        Parse error files and relate relevant information about unaligned files
        '''
        if not self.verbose:
            return
        error_regex = re.compile(r'Did not successfully decode file (\w+),')
        too_little_data_regex = re.compile(
            r'Gaussian has too little data but not removing it because it is the last Gaussian')
        skipped_transition_regex = re.compile(r'(\d+) out of (\d+) transition-states skipped due to insuffient data')

        log_like_regex = re.compile(r'Overall avg like per frame = ([-0-9.]+|nan) over (\d+) frames')
        error_files = []
        for i in range(self.num_jobs):
            path = os.path.join(directory, 'align.{}.{}.log'.format(iteration - 1, i))
            if not os.path.exists(path):
                continue
            with open(path, 'r') as f:
                error_files.extend(error_regex.findall(f.read()))
        update_path = os.path.join(directory, 'update.{}.log'.format(iteration))
        if os.path.exists(update_path):
            with open(update_path, 'r') as f:
                data = f.read()
                m = log_like_regex.search(data)
                if m is not None:
                    log_like, tot_frames = m.groups()
                    if log_like == 'nan':
                        raise (NoSuccessfulAlignments('Could not align any files.  Too little data?'))
                    self.call_back('log-likelihood', float(log_like))
                skipped_transitions = skipped_transition_regex.search(data)
                self.call_back('skipped transitions', *skipped_transitions.groups())
                num_too_little_data = len(too_little_data_regex.findall(data))
                self.call_back('missing data gaussians', num_too_little_data)
        if error_files:
            self.call_back('could not align', error_files)

    def _align_fmllr(self):
        '''
        Align the dataset using speaker-adapted transforms
        '''
        model_directory = self.tri_directory        # Get final.mdl from here
        first_output_directory = self.tri_ali_directory
        second_output_directory = self.tri_fmllr_ali_directory
        self._align_si(fmllr=False)
        sil_phones = self.dictionary.silence_csl

        log_dir = os.path.join(first_output_directory, 'log')
        os.makedirs(log_dir, exist_ok=True)

        calc_fmllr(first_output_directory, self.corpus.split_directory,
                   sil_phones, self.num_jobs, self.tri_fmllr_config, initial=True)
        optional_silence = self.dictionary.optional_silence_csl
        align(0, first_output_directory, self.corpus.split_directory,
              optional_silence, self.num_jobs, self.tri_fmllr_config)

        # Copy into the "correct" tri_fmllr_ali output directory
        for file in glob.glob(os.path.join(first_output_directory, 'ali.*')):
            shutil.copy(file, second_output_directory)
        shutil.copy(os.path.join(first_output_directory, 'tree'), second_output_directory)
        shutil.copy(os.path.join(first_output_directory, 'final.mdl'), second_output_directory)

    def _init_tri(self, fmllr=False):
        if fmllr:
            config = self.tri_fmllr_config
            directory = self.tri_fmllr_directory
            #align_directory = self.tri_fmllr_ali_directory
            align_directory = self.tri_ali_directory
        else:
            config = self.tri_config
            directory = self.tri_directory
            align_directory = self.mono_ali_directory
        # N.B.: Commented out for development
        if os.path.exists(os.path.join(directory, '1.mdl')):
            return
        if fmllr:
            print('Initializing speaker-adapted triphone training...')
        else:
            print('Initializing triphone training...')
        context_opts = []
        ci_phones = self.dictionary.silence_csl

        tree_stats(directory, align_directory,
                   self.corpus.split_directory, ci_phones, self.num_jobs)
        log_path = os.path.join(directory, 'log', 'questions.log')
        tree_path = os.path.join(directory, 'tree')
        treeacc_path = os.path.join(directory, 'treeacc')
        sets_int_path = os.path.join(self.dictionary.phones_dir, 'sets.int')
        roots_int_path = os.path.join(self.dictionary.phones_dir, 'roots.int')
        extra_question_int_path = os.path.join(self.dictionary.phones_dir, 'extra_questions.int')
        topo_path = os.path.join(self.dictionary.output_directory, 'topo')
        questions_path = os.path.join(directory, 'questions.int')
        questions_qst_path = os.path.join(directory, 'questions.qst')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('cluster-phones')] + context_opts +
                            [treeacc_path, sets_int_path, questions_path], stderr=logf)

        with open(extra_question_int_path, 'r') as inf, \
                open(questions_path, 'a') as outf:
            for line in inf:
                outf.write(line)

        log_path = os.path.join(directory, 'log', 'compile_questions.log')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('compile-questions')] + context_opts +
                            [topo_path, questions_path, questions_qst_path],
                            stderr=logf)

        log_path = os.path.join(directory, 'log', 'build_tree.log')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('build-tree')] + context_opts +
                            ['--verbose=1', '--max-leaves={}'.format(config.initial_gauss_count),
                             '--cluster-thresh={}'.format(config.cluster_threshold),
                             treeacc_path, roots_int_path, questions_qst_path,
                             topo_path, tree_path], stderr=logf)

        log_path = os.path.join(directory, 'log', 'init_model.log')
        occs_path = os.path.join(directory, '0.occs')
        mdl_path = os.path.join(directory, '0.mdl')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('gmm-init-model'),
                             '--write-occs=' + occs_path, tree_path, treeacc_path,
                             topo_path, mdl_path], stderr=logf)

        log_path = os.path.join(directory, 'log', 'mixup.log')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('gmm-mixup'),
                             '--mix-up={}'.format(config.initial_gauss_count),
                             mdl_path, occs_path, mdl_path], stderr=logf)

        #os.remove(treeacc_path)

        compile_train_graphs(directory, self.dictionary.output_directory,
                             self.corpus.split_directory, self.num_jobs)

        # Debugging
        #print("@@@@", os.path.exists(os.path.join(directory, '0.mdl')))
        #print("@@@@", os.path.exists(os.path.join(directory, '0.occs')))

        shutil.copy(occs_path, os.path.join(directory, '1.occs'))
        shutil.copy(mdl_path, os.path.join(directory, '1.mdl'))

        convert_alignments(directory, align_directory, self.num_jobs)

        if os.path.exists(os.path.join(align_directory, 'trans.0')):
            for i in range(self.num_jobs):
                shutil.copy(os.path.join(align_directory, 'trans.{}'.format(i)),
                            os.path.join(directory, 'trans.{}'.format(i)))

    def train_tri_fmllr(self):
        '''
        Perform speaker-adapted triphone training
        '''
        # Commented out for testing
        if os.path.exists(self.tri_fmllr_final_model_path):
            print('Triphone FMLLR training already done, using previous final.mdl')
            return

        if not os.path.exists(self.tri_ali_directory):
            self._align_fmllr()

        #self._align_fmllr()

        os.makedirs(os.path.join(self.tri_fmllr_directory, 'log'), exist_ok=True)
        self._init_tri(fmllr=True)
        self._do_tri_fmllr_training()

    def _do_tri_fmllr_training(self):
        self.call_back('Beginning speaker-adapted triphone training...')
        self._do_training(self.tri_fmllr_directory, self.tri_fmllr_config)

    def _do_training(self, directory, config):
        if config.realign_iters is None:
            config.realign_iters = list(range(0, config.num_iters, 10))
        num_gauss = config.initial_gauss_count
        sil_phones = self.dictionary.silence_csl
        inc_gauss = config.inc_gauss_count
        if self.call_back == print:
            iters = tqdm(range(1, config.num_iters))
        else:
            iters = range(1, config.num_iters)
        log_directory = os.path.join(directory, 'log')
        for i in iters:
            model_path = os.path.join(directory, '{}.mdl'.format(i))
            occs_path = os.path.join(directory, '{}.occs'.format(i + 1))
            next_model_path = os.path.join(directory, '{}.mdl'.format(i + 1))
            # N.B.: Commented out for development
            if os.path.exists(next_model_path):
                continue
            if i in config.realign_iters:
                align(i, directory, self.corpus.split_directory,
                      self.dictionary.optional_silence_csl,
                      self.num_jobs, config,
                      feature_name='cmvnsplicetransformfeats')
            if config.do_fmllr and i in config.fmllr_iters:
                calc_fmllr(directory, self.corpus.split_directory, sil_phones,
                           self.num_jobs, config, initial=False, iteration=i)

            if config.do_lda_mllt and i <= config.num_iters:
                calc_lda_mllt(directory, self.corpus.split_directory,   # Could change this to make ali directory later
                #calc_lda_mllt(self.lda_mllt_ali_directory, sil_phones,
                              self.lda_mllt_directory, sil_phones,
                              self.num_jobs, config, config.num_iters,
                              initial=False, iteration=i, corpus=self.corpus)

            acc_stats(i, directory, self.corpus.split_directory, self.num_jobs,
                      config.do_fmllr, do_lda_mllt=config.do_lda_mllt)
            log_path = os.path.join(log_directory, 'update.{}.log'.format(i))
            with open(log_path, 'w') as logf:
                acc_files = [os.path.join(directory, '{}.{}.acc'.format(i, x))
                             for x in range(self.num_jobs)]
                est_proc = subprocess.Popen([thirdparty_binary('gmm-est'),
                                             '--write-occs=' + occs_path,
                                             '--mix-up=' + str(num_gauss), '--power=' + str(config.power),
                                             model_path,
                                             "{} - {}|".format(thirdparty_binary('gmm-sum-accs'),
                                                               ' '.join(map(make_path_safe, acc_files))),
                                             next_model_path],
                                            stderr=logf)
                est_proc.communicate()
            self.parse_log_directory(log_directory, i)
            if i < config.max_iter_inc:
                num_gauss += inc_gauss
        # Debugging
        #print("dir where final.mdl is:", directory)
        shutil.copy(os.path.join(directory, '{}.mdl'.format(config.num_iters)),
                    os.path.join(directory, 'final.mdl'))
        #print("gmm info:")
        #gmm_info_proc = subprocess.Popen([thirdparty_binary('gmm-info'),
        #                                 os.path.join(directory, 'final.mdl')])
        #gmm_info_proc.communicate()

        #print("moving final occs:")
        shutil.copy(os.path.join(directory, '{}.occs'.format(config.num_iters)),
                    os.path.join(directory, 'final.occs'))

        if config.do_lda_mllt:
            #print("moving final mat:")
            shutil.copy(os.path.join(directory, '{}.mat'.format(config.num_iters-1)),
                        os.path.join(directory, 'final.mat'))

    # Big paste (clean up/refactor later)

    def train_lda_mllt(self):
        '''
        Perform LDA + MLLT training
        '''
        # N.B.: Left commented out for development
        if os.path.exists(self.lda_mllt_final_model_path):
            print('LDA + MLLT training already done, using previous final.mdl')
            return

        # N.B: The function _align_lda_mllt() is half-developed, but there doesn't seem to
        # be a reason for it to actually ever be called (since people will always have
        # fmllr done immediately before in the pipeline. Can clean/delete later if determined
        # that we need to actually use it somewhere or not).
        #if not os.path.exists(self.lda_mllt_ali_directory):
        #    self._align_lda_mllt()
        #self._align_lda_mllt()  # half implemented, can come back later or make people run from fmllr

        os.makedirs(os.path.join(self.lda_mllt_directory, 'log'), exist_ok=True)

        self._init_lda_mllt()
        self._do_lda_mllt_training()

    def _init_lda_mllt(self):
        '''
        Initialize LDA + MLLT training.
        '''
        config = self.lda_mllt_config
        directory = self.lda_mllt_directory
        align_directory = self.tri_fmllr_ali_directory  # The previous
        mdl_dir = self.tri_fmllr_directory
        # N.B.: Left commented out for development
        if os.path.exists(os.path.join(directory, '1.mdl')):
            return
        print('Initializing LDA + MLLT training...')

        context_opts = []
        ci_phones = self.dictionary.silence_csl

        log_path = os.path.join(directory, 'log', 'questions.log')
        tree_path = os.path.join(directory, 'tree')
        treeacc_path = os.path.join(directory, 'treeacc')
        sets_int_path = os.path.join(self.dictionary.phones_dir, 'sets.int')
        roots_int_path = os.path.join(self.dictionary.phones_dir, 'roots.int')
        extra_question_int_path = os.path.join(self.dictionary.phones_dir, 'extra_questions.int')
        topo_path = os.path.join(self.dictionary.output_directory, 'topo')
        questions_path = os.path.join(directory, 'questions.int')
        questions_qst_path = os.path.join(directory, 'questions.qst')

        final_mdl_path = os.path.join(self.tri_fmllr_directory)

        # Accumulate LDA stats
        lda_acc_stats(directory, self.corpus.split_directory, align_directory, config, ci_phones, self.num_jobs)

        # Accumulating tree stats
        self.corpus._norm_splice_transform_feats(self.lda_mllt_directory)
        tree_stats(directory, align_directory, self.corpus.split_directory, ci_phones,
                   self.num_jobs, feature_name='cmvnsplicetransformfeats')

        # Getting questions for tree clustering
        log_path = os.path.join(directory, 'log', 'cluster_phones.log')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('cluster-phones')] + context_opts +
                            [treeacc_path, sets_int_path, questions_path], stderr=logf)

        with open(extra_question_int_path, 'r') as inf, \
                open(questions_path, 'a') as outf:
            for line in inf:
                outf.write(line)

        log_path = os.path.join(directory, 'log', 'compile_questions.log')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('compile-questions')] + context_opts +
                            [topo_path, questions_path, questions_qst_path],
                            stderr=logf)

        # Building the tree
        log_path = os.path.join(directory, 'log', 'build_tree.log')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('build-tree')] + context_opts +
                            ['--verbose=1', '--max-leaves={}'.format(config.initial_gauss_count),
                             '--cluster-thresh={}'.format(config.cluster_threshold),
                             treeacc_path, roots_int_path, questions_qst_path,
                             topo_path, tree_path], stderr=logf)

        # Initializing the model
        log_path = os.path.join(directory, 'log', 'init_model.log')
        occs_path = os.path.join(directory, '0.occs')
        mdl_path = os.path.join(directory, '0.mdl')
        with open(log_path, 'w') as logf:
            subprocess.call([thirdparty_binary('gmm-init-model'),
                             '--write-occs=' + occs_path, tree_path, treeacc_path,
                             topo_path, mdl_path], stderr=logf)
        shutil.copy(mdl_path, os.path.join(directory, '1.mdl'))
        shutil.copy(occs_path, os.path.join(directory, '1.occs'))

        convert_alignments(directory, align_directory, self.num_jobs)

        compile_train_graphs(directory, self.dictionary.output_directory,
                             self.corpus.split_directory, self.num_jobs)

        if os.path.exists(os.path.join(align_directory, 'trans.0')):            # ?
            for i in range(self.num_jobs):                                      # ?
                shutil.copy(os.path.join(align_directory, 'trans.{}'.format(i)),# ?
                            os.path.join(directory, 'trans.{}'.format(i)))

    def _do_lda_mllt_training(self):
        self.call_back('Beginning LDA + MLLT training...')
        self._do_training(self.lda_mllt_directory, self.lda_mllt_config)

    def train_nnet_basic(self):
        os.makedirs(os.path.join(self.nnet_basic_directory, 'log'), exist_ok=True)

        split_directory = self.corpus.split_directory
        config = self.nnet_basic_config
        tri_fmllr_config = self.tri_fmllr_config
        directory = self.nnet_basic_directory
        nnet_align_directory = self.nnet_basic_ali_directory
        align_directory = self.tri_fmllr_ali_directory
        lda_directory = self.lda_mllt_directory
        egs_directory = os.path.join(directory, 'egs')
        training_directory = self.corpus.output_directory

        sets_int_path = os.path.join(self.dictionary.phones_dir, 'sets.int')
        roots_int_path = os.path.join(self.dictionary.phones_dir, 'roots.int')
        extra_question_int_path = os.path.join(self.dictionary.phones_dir, 'extra_questions.int')
        topo_path = os.path.join(self.dictionary.output_directory, 'topo')
        L_fst_path = os.path.join(self.dictionary.output_directory, 'L.fst')
        ali_tree_path = os.path.join(align_directory, 'tree')
        shutil.copy(ali_tree_path, os.path.join(directory, 'tree'))

        mdl_path = os.path.join(align_directory, 'final.mdl')
        raw_feats = os.path.join(training_directory, 'feats.scp')

        tree_info_proc = subprocess.Popen([thirdparty_binary('tree-info'),
                                          os.path.join(align_directory, 'tree')],
                                          stdout=subprocess.PIPE)
        tree_info = tree_info_proc.stdout.read()
        tree_info = tree_info.split()
        num_leaves = tree_info[1]
        num_leaves = num_leaves.decode("utf-8")

        lda_dim = 40 # Hard coded, could paramaterize this/make safer, but it's always 40 for LDA

        # Extract iVectors
        self._extract_ivectors()

        # Get LDA matrix
        fixed_ivector_dir = self.extracted_ivector_directory
        #print("FIXED IVECTOR DIR:", fixed_ivector_dir)
        get_lda_nnet(directory, align_directory, fixed_ivector_dir, training_directory,
                     split_directory, raw_feats, self.dictionary.optional_silence_csl, config, self.num_jobs)

        log_path = os.path.join(directory, 'log', 'lda_matrix.log')
        with open(log_path, 'w') as logf:
            acc_files = [os.path.join(directory, 'lda.{}.acc'.format(x))
                         for x in range(self.num_jobs)]
            sum_lda_accs_proc = subprocess.Popen([thirdparty_binary('sum-lda-accs'),
                                                 os.path.join(directory, 'lda.acc')]
                                                 + acc_files,
                                                 stderr=logf)
            sum_lda_accs_proc.communicate()

            lda_mat_proc = subprocess.Popen([thirdparty_binary('nnet-get-feature-transform'),
                                            '--dim=' + str(lda_dim),
                                            os.path.join(directory, 'lda.mat'),
                                            os.path.join(directory, 'lda.acc')],
                                            stderr=logf)
            lda_mat_proc.communicate()
        lda_mat_path = os.path.join(directory, 'lda.mat')

        # Get examples for training
        os.makedirs(egs_directory, exist_ok=True)

        # # Get valid uttlist and train subset uttlist
        # (same issue with hacky paths throughout - these scripts need to be rewritten
        # in python for any release.)

        shuffle_list_path = "/Users/mlml/Documents/Project/kaldi2/egs/wsj/s5/utils/shuffle_list.pl"
        #shuffle_list_path = "/data/acoles/acoles/kaldi/egs/wsj/s5/utils/shuffle_list.pl"
        filter_scp_path = "/Users/mlml/Documents/Project/kaldi2/egs/wsj/s5/utils/filter_scp.pl"
        #filter_scp_path = "/data/acoles/acoles/kaldi/egs/wsj/s5/utils/filter_scp.pl"
        valid_uttlist = os.path.join(directory, 'valid_uttlist')
        train_subset_uttlist = os.path.join(directory, 'train_subset_uttlist')
        training_feats = os.path.join(directory, 'nnet_training_feats')
        num_utts_subset = 300
        log_path = os.path.join(directory, 'log', 'training_egs_feats.log')
        with open(log_path, 'w') as logf:
            with open(valid_uttlist, 'w') as outf:
                valid_uttlist_proc = subprocess.Popen(['awk', '{print $1}',
                                                      os.path.join(training_directory, 'utt2spk')],
                                                      stdout=subprocess.PIPE,
                                                      stderr=logf)
                shuffle_list_proc = subprocess.Popen([shuffle_list_path],
                                                     stdin=valid_uttlist_proc.stdout,
                                                     stdout=subprocess.PIPE,
                                                     stderr=logf)
                head_proc = subprocess.Popen(['head', '-{}'.format(num_utts_subset)],
                                             stdin=shuffle_list_proc.stdout,
                                             stdout=outf,
                                             stderr=logf)
                head_proc.communicate()

            with open(train_subset_uttlist, 'w') as outf:
                awk_proc = subprocess.Popen(['awk', '{print $1}',
                                             os.path.join(training_directory, 'utt2spk')],
                                             stdout=subprocess.PIPE,
                                             stderr=logf)
                filter_scp_proc = subprocess.Popen([filter_scp_path,
                                                   '--exclude', valid_uttlist],
                                                   stdin=awk_proc.stdout,
                                                   stdout=subprocess.PIPE,
                                                   stderr=logf)
                shuffle_list_proc = subprocess.Popen([shuffle_list_path],
                                                     stdin=filter_scp_proc.stdout,
                                                     stdout=subprocess.PIPE,
                                                     stderr=logf)
                head_proc = subprocess.Popen(['head', '-{}'.format(num_utts_subset)],
                                             stdin=shuffle_list_proc.stdout,
                                             stdout=outf,
                                             stderr=logf)
                head_proc.communicate()

        get_egs(directory, egs_directory, training_directory, split_directory, align_directory,
                fixed_ivector_dir, training_feats, valid_uttlist,
                train_subset_uttlist, config, self.num_jobs)

        """# Combine all examples (could integrate validation diagnostics, etc., later-- see egs functions)
        egs_files = []
        for file in os.listdir(egs_directory):
            if file.startswith('egs'):
                egs_files.append(file)
        with open(os.path.join(egs_directory, 'all_egs.egs'), 'wb') as outfile:
            for egs_file in egs_files:
                with open(os.path.join(egs_directory, egs_file), 'rb') as infile:
                    for line in infile:
                        outfile.write(line)"""

        # Initialize neural net
        #stddev = float(1.0/config.hidden_layer_dim**0.5)
        stddev = float(1.0/config.pnorm_input_dim**0.5)
        online_preconditioning_opts = 'alpha={} num-samples-history={} update-period={} rank-in={} rank-out={} max-change-per-sample={}'.format(config.alpha, config.num_samples_history, config.update_period, config.precondition_rank_in, config.precondition_rank_out, config.max_change_per_sample)
        nnet_config_path = os.path.join(directory, 'nnet.config')
        hidden_config_path = os.path.join(directory, 'hidden.config')
        ivector_dim_path = os.path.join(directory, 'ivector_dim')
        with open(ivector_dim_path, 'r') as inf:
            ivector_dim = inf.read().strip()
        #print("here is the ivector dim:", ivector_dim)
        feat_dim = 13 + int(ivector_dim)

        with open(nnet_config_path, 'w') as nc:
            #print("feat dim:", feat_dim)
            nc.write('SpliceComponent input-dim={} left-context={} right-context={} const-component-dim={}\n'.format(feat_dim, config.splice_width, config.splice_width, ivector_dim))
            nc.write('FixedAffineComponent matrix={}\n'.format(lda_mat_path))
            nc.write('AffineComponentPreconditionedOnline input-dim={} output-dim={} {} learning-rate={} param-stddev={} bias-stddev={}\n'.format(lda_dim, config.pnorm_input_dim, online_preconditioning_opts, config.initial_learning_rate, stddev, config.bias_stddev))
            #nc.write('TanhComponent dim={}\n'.format(config.hidden_layer_dim))
            nc.write('PnormComponent input-dim={} output-dim={} p={}\n'.format(config.pnorm_input_dim, config.pnorm_output_dim, config.p))
            nc.write('NormalizeComponent dim={}\n'.format(config.pnorm_output_dim))
            #nc.write('AffineComponent input-dim={} output-dim={} learning-rate={} param-stddev={} bias-stddev={}\n'.format(config.hidden_layer_dim, num_leaves, config.initial_learning_rate, stddev, config.bias_stddev))
            nc.write('AffineComponentPreconditionedOnline input-dim={} output-dim={} {} learning-rate={} param-stddev={} bias-stddev={}\n'.format(config.pnorm_output_dim, num_leaves, online_preconditioning_opts, config.initial_learning_rate, stddev, config.bias_stddev))
            nc.write('SoftmaxComponent dim={}\n'.format(num_leaves))

        with open(hidden_config_path, 'w') as nc:
            #nc.write('AffineComponent input-dim={} output-dim={} learning-rate={} param-stddev={} bias-stddev={}\n'.format(config.hidden_layer_dim, config.hidden_layer_dim, config.initial_learning_rate, stddev, config.bias_stddev))
            nc.write('AffineComponentPreconditionedOnline input-dim={} output-dim={} {} learning-rate={} param-stddev={} bias-stddev={}\n'.format(config.pnorm_output_dim, config.pnorm_input_dim, online_preconditioning_opts, config.initial_learning_rate, stddev, config.bias_stddev))
            #nc.write('TanhComponent dim={}\n'.format(config.hidden_layer_dim))
            nc.write('PnormComponent input-dim={} output-dim={} p={}\n'.format(config.pnorm_input_dim, config.pnorm_output_dim, config.p))
            nc.write('NormalizeComponent dim={}\n'.format(config.pnorm_output_dim))

        log_path = os.path.join(directory, 'log', 'nnet_init.log')
        with open(log_path, 'w') as logf:
            nnet_am_init_proc = subprocess.Popen([thirdparty_binary('nnet-am-init'),
                                                 os.path.join(align_directory, 'tree'),
                                                 topo_path,
                                                 "{} {} -|".format(thirdparty_binary('nnet-init'),
                                                                   nnet_config_path),
                                                os.path.join(directory, '0.mdl')],
                                                stderr=logf)
            nnet_am_init_proc.communicate()

            nnet_am_info = subprocess.Popen([thirdparty_binary('nnet-am-info'),
                                            os.path.join(directory, '0.mdl')],
                                            stderr=logf)
            nnet_am_info.communicate()

        # Train transition probabilities and set priors
        #   First combine all previous alignments
        ali_files = glob.glob(os.path.join(align_directory, 'ali.*'))
        prev_ali_path = os.path.join(directory, 'prev_ali.')
        with open(prev_ali_path, 'wb') as outfile:
            for ali_file in ali_files:
                with open(os.path.join(align_directory, ali_file), 'rb') as infile:
                    for line in infile:
                        outfile.write(line)
        nnet_train_trans(directory, align_directory, prev_ali_path, self.num_jobs)

        # Get iteration at which we will mix up
        num_tot_iters = config.num_epochs * config.iters_per_epoch
        finish_add_layers_iter = config.num_hidden_layers * config.add_layers_period
        first_modify_iter = finish_add_layers_iter + config.add_layers_period
        mix_up_iter = (num_tot_iters + finish_add_layers_iter)/2

        # Get iterations at which we will realign
        realign_iters = []
        if config.realign_times != 0:
            div = config.realign_times + 1 # (e.g. realign 2 times = iterations split into 3 sets)
            split_iters = np.array_split(range(num_tot_iters), div)
            for group in split_iters:
                realign_iters.append(group[0])

        # Training loop
        for i in range(num_tot_iters):
            print("On iteration {}...".format(i))
            model_path = os.path.join(directory, '{}.mdl'.format(i))
            next_model_path = os.path.join(directory, '{}.mdl'.format(i + 1))

            # Combine all examples (could integrate validation diagnostics, etc., later-- see egs functions)
            egs_files = []
            for file in os.listdir(egs_directory):
                if file.startswith('egs'):
                    egs_files.append(file)
            with open(os.path.join(egs_directory, 'all_egs.egs'), 'wb') as outfile:
                for egs_file in egs_files:
                    with open(os.path.join(egs_directory, egs_file), 'rb') as infile:
                        for line in infile:
                            outfile.write(line)

            # Get accuracy rates for the current iteration (to pull out graph later)
            #compute_prob(i, directory, egs_directory, model_path, self.num_jobs)
            log_path = os.path.join(directory, 'log', 'compute_prob_train.{}.log'.format(i))
            with open(log_path, 'w') as logf:
                compute_prob_proc = subprocess.Popen([thirdparty_binary('nnet-compute-prob'),
                                                     model_path,
                                                     'ark:{}/all_egs.egs'.format(egs_directory)],
                                                     stderr=logf)
                compute_prob_proc.communicate()

            # Pull out and save graphs
            #get_accuracy_graph(os.path.join(directory, 'log'), os.path.join(directory, 'log'))

            # If it is NOT the first iteration,
            # AND we still have layers to add,
            # AND it's the right time to add a layer...
            if i > 0 and i <= ((config.num_hidden_layers-1)*config.add_layers_period) and ((i-1)%config.add_layers_period) == 0:
                # Add a new hidden layer
                mdl = os.path.join(directory, 'tmp{}.mdl'.format(i))
                log_path = os.path.join(directory, 'log', 'temp_mdl.{}.log'.format(i))
                with open(log_path, 'w') as logf:
                    with open(mdl, 'w') as outf:
                        tmp_mdl_init_proc = subprocess.Popen([thirdparty_binary('nnet-init'),
                                                            '--srand={}'.format(i),
                                                            os.path.join(directory, 'hidden.config'),
                                                            '-'],
                                                            stdout=subprocess.PIPE,
                                                            stderr=logf)
                        tmp_mdl_ins_proc = subprocess.Popen([thirdparty_binary('nnet-insert'),
                                                            os.path.join(directory, '{}.mdl'.format(i)),
                                                            '-', '-'],
                                                            stdin=tmp_mdl_init_proc.stdout,
                                                            stdout=outf,
                                                            stderr=logf)
                        tmp_mdl_ins_proc.communicate()

            # Otherwise just use the past model
            else:
                mdl = os.path.join(directory, '{}.mdl'.format(i))

            # Shuffle examples and train nets with SGD
            nnet_train(directory, egs_directory, mdl, i, self.num_jobs)

            # Get nnet list from the various jobs on this iteration
            nnets_list = [os.path.join(directory, '{}.{}.mdl'.format((i+1), x))
                         for x in range(self.num_jobs)]

            if (i+1) >= num_tot_iters:
                learning_rate = config.final_learning_rate
            else:
                learning_rate = config.initial_learning_rate * math.exp(i * math.log(config.final_learning_rate/config.initial_learning_rate)/num_tot_iters)

            log_path = os.path.join(directory, 'log', 'average.{}.log'.format(i))
            with open(log_path, 'w') as logf:
                nnet_avg_proc = subprocess.Popen([thirdparty_binary('nnet-am-average')]
                                                 + nnets_list
                                                 + ['-'],
                                                 stdout=subprocess.PIPE,
                                                 stderr=logf)
                nnet_copy_proc = subprocess.Popen([thirdparty_binary('nnet-am-copy'),
                                                  '--learning-rate={}'.format(learning_rate),
                                                  '-',
                                                  os.path.join(directory, '{}.mdl'.format(i+1))],
                                                  stdin=nnet_avg_proc.stdout,
                                                  stderr=logf)
                nnet_copy_proc.communicate()

            # If it's the right time, do mixing up
            if config.mix_up > 0 and i == mix_up_iter:
                log_path = os.path.join(directory, 'log', 'mix_up.{}.log'.format(i))
                with open(log_path, 'w') as logf:
                    nnet_am_mixup_proc = subprocess.Popen([thirdparty_binary('nnet-am-mixup'),
                                                          '--min-count=10',
                                                          '--num-mixtures={}'.format(config.mix_up),
                                                          os.path.join(directory, '{}.mdl'.format(i+1)),
                                                          os.path.join(directory, '{}.mdl'.format(i+1))],
                                                          stderr=logf)
                    nnet_am_mixup_proc.communicate()

            # Realign if it's the right time
            if i in realign_iters:
                prev_egs_directory = egs_directory
                egs_directory = os.path.join(directory, 'egs_{}'.format(i))
                os.makedirs(egs_directory, exist_ok=True)

                #   Get average posterior for purposes of adjusting priors
                get_average_posteriors(i, directory, prev_egs_directory, config, self.num_jobs)
                log_path = os.path.join(directory, 'log', 'vector_sum_exterior.{}.log'.format(i))
                vectors_to_sum = glob.glob(os.path.join(directory, 'post.{}.*.vec'.format(i)))
                #vectors = glob.glob(os.path.join(directory, 'post.{}.*.vec'.format(i)))
                #vectors_to_sum_path = os.path.join(directory, 'vectors_to_sum.')
                #with open(vectors_to_sum_path, 'wb') as outfile:
                #    for vector in vectors:
                #        with open(os.path.join(directory, vector), 'rb') as infile:
                #            for line in infile:
                #                outfile.write(line)

                with open(log_path, 'w') as logf:
                    vector_sum_proc = subprocess.Popen([thirdparty_binary('vector-sum')]
                                                       + vectors_to_sum
                                                       + [os.path.join(directory, 'post.{}.vec'.format(i))
                                                       #os.path.join(directory, '{}.mdl'.format(i))],
                                                       #vectors_to_sum_path,
                                                       #os.path.join(directory, 'post.{}.vec'.format(i))
                                                       ],
                                                       stderr=logf)
                    vector_sum_proc.communicate()

                #   Readjust priors based on computed posteriors
                log_path = os.path.join(directory, 'log', 'adjust_priors.{}.log'.format(i))
                with open(log_path, 'w') as logf:
                    nnet_adjust_priors_proc = subprocess.Popen([thirdparty_binary('nnet-adjust-priors'),
                                                               os.path.join(directory, '{}.mdl'.format(i)),
                                                               os.path.join(directory, 'post.{}.vec'.format(i)),
                                                               os.path.join(directory, '{}.mdl'.format(i))],
                                                               stderr=logf)
                    nnet_adjust_priors_proc.communicate()

                #   Realign:
                #       Compile train graphs (gets fsts.{} for alignment)
                compile_train_graphs(directory, self.dictionary.output_directory,
                                     self.corpus.split_directory, self.num_jobs)

                #       Get alignment feats
                nnet_get_align_feats(directory, self.corpus.split_directory, fixed_ivector_dir, config, self.num_jobs)

                #       Do alignment
                nnet_align(i, directory,
                      self.dictionary.optional_silence_csl,
                      self.num_jobs, config)

                #     Finally, relabel the egs
                ali_files = glob.glob(os.path.join(directory, 'ali.*'))
                alignments = os.path.join(directory, 'alignments.')
                with open(alignments, 'wb') as outfile:
                    for ali_file in ali_files:
                        with open(os.path.join(directory, ali_file), 'rb') as infile:
                            for line in infile:
                                outfile.write(line)
                relabel_egs(i, directory, prev_egs_directory, alignments, egs_directory, self.num_jobs)

        # Rename the final model
        shutil.copy(os.path.join(directory, '{}.mdl'.format(num_tot_iters-1)), os.path.join(directory, 'final.mdl'))

        # Compile train graphs (gets fsts.{} for alignment)
        compile_train_graphs(directory, self.dictionary.output_directory,
                             self.corpus.split_directory, self.num_jobs)

        # Get alignment feats
        nnet_get_align_feats(directory, self.corpus.split_directory, fixed_ivector_dir, config, self.num_jobs)

        # Do alignment
        nnet_align("final", directory,
              self.dictionary.optional_silence_csl,
              self.num_jobs, config, mdl=os.path.join(directory, 'final.mdl'))

    def _extract_ivectors(self):
        # N.B.: Left commented out for development
        if os.path.exists(self.ivector_extractor_final_model_path):
            print('iVector extractor training already done, using previous final.ie')
            return

        log_dir = os.path.join(self.extracted_ivector_directory, 'log')
        os.makedirs(log_dir, exist_ok=True)

        # N.B.: These paths are hacky and need to be correctly integrated.
        # To do still for release: maybe add arguments to command line to tell MFA which
        # ivector extractor to use, and infer the diag_ubm and lda_mllt paths from that.

        directory = self.extracted_ivector_directory
        #ivector_extractor_dir = "/data/acoles/acoles/Montreal-Forced-Aligner/ivector_extractors/ls_medium/ivector_extractor_ls_medium"
        #ivector_extractor_dir = "../../ivector_extractor"
        ivector_extractor_dir = "/Users/mlml/Documents/GitHub/Montreal-Forced-Aligner/ivector_extractor"
        #diag_ubm_path = "/data/acoles/acoles/Montreal-Forced-Aligner/ivector_extractors/ls_medium/diag_ubm"
        #diag_ubm_path = "../../diag_ubm"
        diag_ubm_path = "/Users/mlml/Documents/GitHub/Montreal-Forced-Aligner/diag_ubm"
        #lda_mllt_path = "/data/acoles/acoles/Montreal-Forced-Aligner/ivector_extractors/ls_medium/lda_mllt"
        #lda_mllt_path = "../../lda_mllt"
        lda_mllt_path = "/Users/mlml/Documents/GitHub/Montreal-Forced-Aligner/lda_mllt"
        split_dir = self.corpus.split_directory
        train_dir = self.corpus.output_directory
        config = self.ivector_extractor_config
        training_directory = self.corpus.output_directory

        # Need to make a directory for corpus with just 2 utterances per speaker
        # (left commented out in case we ever decide to do this)
        #max2_dir = os.path.join(directory, 'max2')
        #os.makedirs(max2_dir, exist_ok=True)
        #mfa_working_dir = os.getcwd()
        #os.chdir("/Users/mlml/Documents/Project/kaldi2/egs/wsj/s5/steps/online/nnet2")
        #opy_data_sh = "/Users/mlml/Documents/Project/kaldi2/egs/wsj/s5/steps/online/nnet2/copy_data_dir.sh"
        #log_path = os.path.join(directory, 'log', 'max2.log')
        #with open(log_path, 'w') as logf:
        #    command = [copy_data_sh, '--utts-per-spk-max', '2', train_dir, max2_dir]
        #    max2_proc = subprocess.Popen(command,
        #                                 stderr=logf)
        #    max2_proc.communicate()
        #os.chdir(mfa_working_dir)

        # Write a "cmvn config" file (this is blank in the actual kaldi code, but it needs the argument passed)
        cmvn_config = os.path.join(directory, 'online_cmvn.conf')
        with open(cmvn_config, 'w') as cconf:
            cconf.write("")

        # Write a "splice config" file
        splice_config = os.path.join(directory, 'splice.conf')
        with open(splice_config, 'w') as sconf:
            sconf.write(config.splice_opts[0])
            sconf.write('\n')
            sconf.write(config.splice_opts[1])

        # Write a "config" file to input to the extraction binary
        ext_config = os.path.join(directory, 'ivector_extractor.conf')
        with open(ext_config, 'w') as ieconf:
            ieconf.write('--cmvn-config={}\n'.format(cmvn_config))
            ieconf.write('--ivector-period={}\n'.format(config.ivector_period))
            ieconf.write('--splice-config={}\n'.format(splice_config))
            ieconf.write('--lda-matrix={}\n'.format(os.path.join(lda_mllt_path, 'final.mat')))
            ieconf.write('--global-cmvn-stats={}\n'.format(os.path.join(diag_ubm_path, 'global_cmvn.stats')))
            ieconf.write('--diag-ubm={}\n'.format(os.path.join(diag_ubm_path, 'final.dubm')))
            ieconf.write('--ivector-extractor={}\n'.format(os.path.join(ivector_extractor_dir, 'final.ie')))
            ieconf.write('--num-gselect={}\n'.format(config.num_gselect))
            ieconf.write('--min-post={}\n'.format(config.min_post))
            ieconf.write('--posterior-scale={}\n'.format(config.posterior_scale))
            ieconf.write('--max-remembered-frames=1000\n')
            ieconf.write('--max-count={}\n'.format(0))

        # Extract iVectors
        extract_ivectors(directory, training_directory, ext_config, config, self.num_jobs)

        # Combine iVectors across jobs
        file_list = []
        for j in range(self.num_jobs):
            file_list.append(os.path.join(directory, 'ivector_online.{}.scp'.format(j)))
        #print("file list:", file_list)

        with open(os.path.join(directory, 'ivector_online.scp'), 'w') as outfile:
            for fname in file_list:
                with open(fname) as infile:
                    for line in infile:
                        outfile.write(line)

max_stars_repo_name: a-coles/Montreal-Forced-Aligner

max_issues_repo_name: a-coles/Montreal-Forced-Aligner

max_forks_repo_name: a-coles/Montreal-Forced-Aligner


text:
#!/usr/bin/env python

import sys

import rospy
from std_msgs.msg import String
from sensor_msgs.msg import Image
from geometry_msgs.msg import Pose, PoseStamped, Point, Point32, PolygonStamped
from cv_bridge import CvBridge, CvBridgeError

from scipy import ndimage, signal
import argparse
import os
import sys
import math
import numpy as np
import copy
import transforms3d as tf3d
import json
import copy

import keras
import tensorflow as tf
import open3d
import ros_numpy

#print(sys.path)
#sys.path.remove('/opt/ros/melodic/lib/python2.7/dist-packages')
import cv2
from PIL import Image as Pilimage

sys.path.append("/RGBDPose")
from RGBDPose import models
from RGBDPose.utils.config import read_config_file, parse_anchor_parameters
from RGBDPose.utils.eval import evaluate
from RGBDPose.utils.keras_version import check_keras_version
from RGBDPose.utils import ply_loader

from object_detector_msgs.srv import get_poses, get_posesResponse
from object_detector_msgs.msg import PoseWithConfidence
from geometry_msgs.msg import PoseArray, Pose
###################################
##### Global Variable Space #######
######## aka. death zone ##########
###################################

# LineMOD
#fxkin = 572.41140
#fykin = 573.57043
#cxkin = 325.26110
#cykin = 242.04899

# YCB-video
#fxkin = 1066.778
#fykin = 1067.487
#cxkin = 312.9869
#cykin = 241.3109

# our Kinect
#fxkin = 575.81573
#fykin = 575.81753
#cxkin = 314.5
#cykin = 235.5

# HSRB
# fxkin = 538.391033
# fykin = 538.085452
# cxkin = 315.30747
# cykin = 233.048356

# magic intrinsics
fxkin = 1066.778
fykin = 1067.487
cxkin = 320.0
cykin = 240.0

def get_evaluation_kiru(pcd_temp_,pcd_scene_,inlier_thres,tf,final_th, model_dia):#queue
    tf_pcd =np.eye(4)
    pcd_temp_.transform(tf)

    mean_temp = np.mean(np.array(pcd_temp_.points)[:, 2])
    mean_scene = np.median(np.array(pcd_scene_.points)[:, 2])
    pcd_diff = mean_scene - mean_temp

    # align model with median depth of scene
    new_pcd_trans = []
    for i, point in enumerate(pcd_temp_.points):
        poi = np.asarray(point)
        poi = poi + [0.0, 0.0, pcd_diff]
        new_pcd_trans.append(poi)
    tf = np.array(tf)
    tf[2, 3] = tf[2, 3] + pcd_diff
    pcd_temp_.points = open3d.Vector3dVector(np.asarray(new_pcd_trans))
    open3d.estimate_normals(pcd_temp_, search_param=open3d.KDTreeSearchParamHybrid(
        radius=5.0, max_nn=10))

    pcd_min = mean_scene - (model_dia * 2)
    pcd_max = mean_scene + (model_dia * 2)
    new_pcd_scene = []
    for i, point in enumerate(pcd_scene_.points):
        if point[2] > pcd_min or point[2] < pcd_max:
            new_pcd_scene.append(point)
    pcd_scene_.points = open3d.Vector3dVector(np.asarray(new_pcd_scene))
    open3d.estimate_normals(pcd_scene_, search_param=open3d.KDTreeSearchParamHybrid(
        radius=5.0, max_nn=10))

    reg_p2p = open3d.registration.registration_icp(pcd_temp_,pcd_scene_ , inlier_thres, np.eye(4),
                                                   open3d.registration.TransformationEstimationPointToPoint(),
                                                   open3d.registration.ICPConvergenceCriteria(max_iteration = 5)) #5?
    tf = np.matmul(reg_p2p.transformation,tf)
    tf_pcd = np.matmul(reg_p2p.transformation,tf_pcd)
    pcd_temp_.transform(reg_p2p.transformation)

    #open3d.estimate_normals(pcd_temp_, search_param=open3d.KDTreeSearchParamHybrid(
    #    radius=2.0, max_nn=30))
    points_unfiltered = np.asarray(pcd_temp_.points)
    last_pcd_temp = []
    for i, normal in enumerate(pcd_temp_.normals):
        if normal[2] < 0:
            last_pcd_temp.append(points_unfiltered[i, :])

    pcd_temp_.points = open3d.Vector3dVector(np.asarray(last_pcd_temp))
    open3d.estimate_normals(pcd_temp_, search_param=open3d.KDTreeSearchParamHybrid(
        radius=5.0, max_nn=30))

    hyper_tresh = inlier_thres
    for i in range(4):
        inlier_thres = reg_p2p.inlier_rmse*2
        hyper_thres = hyper_tresh * 0.75
        if inlier_thres < 1.0:
            inlier_thres = hyper_tresh * 0.75
            hyper_tresh = inlier_thres
        reg_p2p = open3d.registration.registration_icp(pcd_temp_,pcd_scene_ , inlier_thres, np.eye(4),
                                                       open3d.registration.TransformationEstimationPointToPlane(),
                                                       open3d.registration.ICPConvergenceCriteria(max_iteration = 1)) #5?
        tf = np.matmul(reg_p2p.transformation,tf)
        tf_pcd = np.matmul(reg_p2p.transformation,tf_pcd)
        pcd_temp_.transform(reg_p2p.transformation)
    inlier_rmse = reg_p2p.inlier_rmse

    #open3d.draw_geometries([pcd_temp_, pcd_scene_])

    ##Calculate fitness with depth_inlier_th
    if(final_th>0):

        inlier_thres = final_th #depth_inlier_th*2 #reg_p2p.inlier_rmse*3
        reg_p2p = open3d.registration.registration_icp(pcd_temp_,pcd_scene_, inlier_thres, np.eye(4),
                                                       open3d.registration.TransformationEstimationPointToPlane(),
                                                       open3d.registration.ICPConvergenceCriteria(max_iteration = 1)) #5?
        tf = np.matmul(reg_p2p.transformation, tf)
        tf_pcd = np.matmul(reg_p2p.transformation, tf_pcd)
        pcd_temp_.transform(reg_p2p.transformation)

    #open3d.draw_geometries([last_pcd_temp_, pcd_scene_])

    if( np.abs(np.linalg.det(tf[:3,:3])-1)>0.001):
        tf[:3,0]=tf[:3,0]/np.linalg.norm(tf[:3,0])
        tf[:3,1]=tf[:3,1]/np.linalg.norm(tf[:3,1])
        tf[:3,2]=tf[:3,2]/np.linalg.norm(tf[:3,2])
    if( np.linalg.det(tf) < 0) :
        tf[:3,2]=-tf[:3,2]

    return tf,inlier_rmse,tf_pcd,reg_p2p.fitness

def create_point_cloud(depth, fx, fy, cx, cy, ds):

    rows, cols = depth.shape

    depRe = depth.reshape(rows * cols)
    zP = np.multiply(depRe, ds)

    x, y = np.meshgrid(np.arange(0, cols, 1), np.arange(0, rows, 1), indexing='xy')
    yP = y.reshape(rows * cols) - cy
    xP = x.reshape(rows * cols) - cx
    yP = np.multiply(yP, zP)
    xP = np.multiply(xP, zP)
    yP = np.divide(yP, fy)
    xP = np.divide(xP, fx)

    cloud_final = np.transpose(np.array((xP, yP, zP)))
    cloud_final[cloud_final[:,2]==0] = np.NaN

    return cloud_final

def preprocess_image(x, mode='caffe'):
    x = x.astype(np.float32)

    if mode == 'tf':
        x /= 127.5
        x -= 1.
    elif mode == 'caffe':
        x[..., 0] -= 103.939
        x[..., 1] -= 116.779
        x[..., 2] -= 123.68

    return x

#################################
############### ROS #############
#################################
class PoseEstimationClass:
    #def __init__(self, model, mesh_path, threshold, topic, graph):
    def __init__(self, model, mesh_path, threshold, topic):
        #event that will block until the info is received
        #attribute for storing the rx'd message
        self._model = model
        self._score_th = threshold
        #self.graph = graph

        self._msg = None
        self.seq = None
        self.time = None
        self.frame_id = None
        self.bridge = CvBridge()
        self.pose_pub = rospy.Publisher("/pyrapose/poses", PoseArray, queue_size=10)
        self.image_sub = rospy.Subscriber(topic, Image, self.callback)
        self.depth_sub = rospy.Subscriber('/hsrb/head_rgbd_sensor/depth_registered/image_raw', Image, self.depth_callback)

        self.threeD_boxes = np.ndarray((22, 8, 3), dtype=np.float32)
        mesh_info = os.path.join(mesh_path, 'models_info.json')
        for key, value in json.load(open(mesh_info)).items():
            fac = 0.001
            x_minus = value['min_x'] * fac
            y_minus = value['min_y'] * fac
            z_minus = value['min_z'] * fac
            x_plus = value['size_x'] * fac + x_minus
            y_plus = value['size_y'] * fac + y_minus
            z_plus = value['size_z'] * fac + z_minus
            three_box_solo = np.array([[x_plus, y_plus, z_plus],
                                   [x_plus, y_plus, z_minus],
                                   [x_plus, y_minus, z_minus],
                                   [x_plus, y_minus, z_plus],
                                   [x_minus, y_plus, z_plus],
                                   [x_minus, y_plus, z_minus],
                                   [x_minus, y_minus, z_minus],
                                   [x_minus, y_minus, z_plus]])
            self.threeD_boxes[int(key), :, :] = three_box_solo
        ply_path = os.path.join(mesh_path, 'obj_000005.ply')
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_6 = open3d.PointCloud()
        self.model_6.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_6 = open3d.PointCloud()
        self.pcd_model_6.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_6, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
        ply_path = mesh_path + '/obj_000008.ply'
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_9 = open3d.PointCloud()
        self.model_9.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_9 = open3d.PointCloud()
        self.pcd_model_9.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_9, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
        ply_path = mesh_path + '/obj_000009.ply'
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_10 = open3d.PointCloud()
        self.model_10.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_10 = open3d.PointCloud()
        self.pcd_model_10.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_10, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
        ply_path = mesh_path + '/obj_000010.ply'
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_11 = open3d.PointCloud()
        self.model_11.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_11 = open3d.PointCloud()
        self.pcd_model_11.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_11, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
        ply_path = mesh_path + '/obj_000021.ply'
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_61 = open3d.PointCloud()
        self.model_61.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_61 = open3d.PointCloud()
        self.pcd_model_61.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_61, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
    def depth_callback(self, data):
        self.depth = data

    def callback(self, data):
        self.seq = data.header.seq
        self.time = data.header.stamp
        self.frame_id = data.header.frame_id
        self._msg = self.bridge.imgmsg_to_cv2(data, "8UC3")
        #self._msg = ros_numpy.numpify(data)
        self._dep =self.bridge.imgmsg_to_cv2(self.depth, "16UC1")

        f_sca_x = 538.391033 / 1066.778
        f_sca_y = 538.085452 / 1067.487
        x_min = 315.30747 * f_sca_x
        x_max = 315.30747 + (640.0 - 315.30747) * f_sca_x
        y_min = 233.04356 * f_sca_y
        y_max = 233.04356 + (480.0 - 233.04356) * f_sca_y
        self._msg = self._msg[int(y_min):int(y_max), int(x_min):int(x_max), :]
        self._msg = cv2.resize(self._msg, (640, 480))

        self._msg = cv2.cvtColor(self._msg, cv2.COLOR_BGR2RGB)
        cv2.imwrite('/stefan/test.png', self._msg)
        self._dep = self._dep[int(y_min):int(y_max), int(x_min):int(x_max)]
        self._dep = cv2.resize(self._dep, (640, 480))

        det_objs, det_poses, det_confs = run_estimation(self._msg, self._dep, self._model, self._score_th, self.threeD_boxes, self.pcd_model_6, self.pcd_model_9, self.pcd_model_10, self.pcd_model_11, self.pcd_model_61)#, self.seq)

        self.publish_pose(det_objs, det_poses, det_confs)
        rospy.sleep(2)

    def publish_pose(self, det_names, det_poses, det_confidences):
        msg = PoseArray()
        msg.header.frame_id = '/head_rgbd_sensor_rgb_frame'
        msg.header.stamp = rospy.Time(0)

        for idx in range(len(det_names)):
            item = Pose()
            item.position.x = det_poses[idx][0]
            item.position.y = det_poses[idx][1]
            item.position.z = det_poses[idx][2]
            item.orientation.w = det_poses[idx][3]
            item.orientation.x = det_poses[idx][4]
            item.orientation.y = det_poses[idx][5]
            item.orientation.z = det_poses[idx][6]
            msg.poses.append(item)
        self.pose_pub.publish(msg)
        # msg = get_posesResponse()
        # for idx in range(len(det_names)):
        #     item = PoseWithConfidence()
        #     item.name = det_names[idx]
        #     item.confidence = det_confidences[idx]
        #     item.pose = Pose()
        #     det_pose = det_poses[idx]
        #     item.pose.position.x = det_pose[0]
        #     item.pose.position.y = det_pose[1]
        #     item.pose.position.z = det_pose[2]
        #     item.pose.orientation.w = det_pose[3]
        #     item.pose.orientation.x = det_pose[4]
        #     item.pose.orientation.y = det_pose[5]
        #     item.pose.orientation.z = det_pose[6]
        #     msg.poses.append(item)

        # self.pose_pub.publish(msg)

class PoseEstimationServer:
    def __init__(self, model, mesh_path, threshold, topic, service_name):
        #event that will block until the info is received
        #attribute for storing the rx'd message
        self._model = model
        self._score_th = threshold

        self._msg = None
        self.seq = None
        self.time = None
        self.frame_id = None
        self.bridge = CvBridge()
        self.topic = topic
        self.pose_pub = rospy.Publisher("/pyrapose/poses", PoseArray, queue_size=10)
        self.pose_srv = rospy.Service(service_name, get_poses, self.callback)
        self.image_sub = rospy.Subscriber(topic, Image, self.image_callback)
        self.depth_sub = rospy.Subscriber('/hsrb/head_rgbd_sensor/depth_registered/image_raw', Image, self.depth_callback)

        self.threeD_boxes = np.ndarray((22, 8, 3), dtype=np.float32)
        mesh_info = os.path.join(mesh_path, 'models_info.json')
        for key, value in json.load(open(mesh_info)).items():
            fac = 0.001
            x_minus = value['min_x'] * fac
            y_minus = value['min_y'] * fac
            z_minus = value['min_z'] * fac
            x_plus = value['size_x'] * fac + x_minus
            y_plus = value['size_y'] * fac + y_minus
            z_plus = value['size_z'] * fac + z_minus
            three_box_solo = np.array([[x_plus, y_plus, z_plus],
                                   [x_plus, y_plus, z_minus],
                                   [x_plus, y_minus, z_minus],
                                   [x_plus, y_minus, z_plus],
                                   [x_minus, y_plus, z_plus],
                                   [x_minus, y_plus, z_minus],
                                   [x_minus, y_minus, z_minus],
                                   [x_minus, y_minus, z_plus]])
            self.threeD_boxes[int(key), :, :] = three_box_solo
        ply_path = os.path.join(mesh_path, 'obj_000005.ply')
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_6 = open3d.PointCloud()
        self.model_6.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_6 = open3d.PointCloud()
        self.pcd_model_6.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_6, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
        ply_path = mesh_path + '/obj_000008.ply'
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_9 = open3d.PointCloud()
        self.model_9.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_9 = open3d.PointCloud()
        self.pcd_model_9.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_9, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
        ply_path = mesh_path + '/obj_000009.ply'
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_10 = open3d.PointCloud()
        self.model_10.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_10 = open3d.PointCloud()
        self.pcd_model_10.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_10, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
        ply_path = mesh_path + '/obj_000010.ply'
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_11 = open3d.PointCloud()
        self.model_11.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_11 = open3d.PointCloud()
        self.pcd_model_11.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_11, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))
        ply_path = mesh_path + '/obj_000021.ply'
        model_vsd = ply_loader.load_ply(ply_path)
        self.model_61 = open3d.PointCloud()
        self.model_61.points = open3d.Vector3dVector(model_vsd['pts'])
        self.pcd_model_61 = open3d.PointCloud()
        self.pcd_model_61.points = open3d.Vector3dVector(model_vsd['pts'])
        open3d.estimate_normals(self.pcd_model_61, search_param=open3d.KDTreeSearchParamHybrid(
        radius=2.0, max_nn=30))

    def image_callback(self, data):
        self.image = data

    def depth_callback(self, data):
        self.depth = data

    def callback(self, req):
        #print(data)
        rospy.wait_for_message(self.topic, Image)
        data = self.image
        self.seq = data.header.seq
        self.time = data.header.stamp
        self.frame_id = data.header.frame_id
        self._msg = self.bridge.imgmsg_to_cv2(data, "8UC3")
        #self._msg = ros_numpy.numpify(data)
        self._dep =self.bridge.imgmsg_to_cv2(self.depth, "16UC1")

        f_sca_x = 538.391033 / 1066.778
        f_sca_y = 538.085452 / 1067.487
        x_min = 315.30747 * f_sca_x
        x_max = 315.30747 + (640.0 - 315.30747) * f_sca_x
        y_min = 233.04356 * f_sca_y
        y_max = 233.04356 + (480.0 - 233.04356) * f_sca_y
        self._msg = self._msg[int(y_min):int(y_max), int(x_min):int(x_max), :]
        self._msg = cv2.resize(self._msg, (640, 480))

        self._msg = cv2.cvtColor(self._msg, cv2.COLOR_BGR2RGB)
        cv2.imwrite('/stefan/test.png', self._msg)
        self._dep = self._dep[int(y_min):int(y_max), int(x_min):int(x_max)]
        self._dep = cv2.resize(self._dep, (640, 480))

        det_objs, det_poses, det_confs = run_estimation(self._msg, self._dep, self._model, self._score_th, self.threeD_boxes, self.pcd_model_6, self.pcd_model_9, self.pcd_model_10, self.pcd_model_11, self.pcd_model_61)#, self.seq)

        msg = self.fill_pose(det_objs, det_poses, det_confs)
        return msg

    def fill_pose(self, det_names, det_poses, det_confidences):
        msg = PoseArray()
        msg.header.frame_id = '/head_rgbd_sensor_rgb_frame'
        msg.header.stamp = rospy.Time(0)

        for idx in range(len(det_names)):
            item = Pose()
            item.position.x = det_poses[idx][0]
            item.position.y = det_poses[idx][1]
            item.position.z = det_poses[idx][2]
            item.orientation.w = det_poses[idx][3]
            item.orientation.x = det_poses[idx][4]
            item.orientation.y = det_poses[idx][5]
            item.orientation.z = det_poses[idx][6]
            msg.poses.append(item)
        self.pose_pub.publish(msg)

        msg = get_posesResponse()
        for idx in range(len(det_names)):
            item = PoseWithConfidence()
            item.name = det_names[idx]
            item.confidence = det_confidences[idx]
            item.pose = Pose()
            det_pose = det_poses[idx]
            item.pose.position.x = det_pose[0]
            item.pose.position.y = det_pose[1]
            item.pose.position.z = det_pose[2]
            item.pose.orientation.w = det_pose[3]
            item.pose.orientation.x = det_pose[4]
            item.pose.orientation.y = det_pose[5]
            item.pose.orientation.z = det_pose[6]
            msg.poses.append(item)

        return msg

#################################
########## RetNetPose ###########
#################################
def get_session():
    """ Construct a modified tf session.
    """
    config = tf.ConfigProto()
    #config.gpu_options.allow_growth = True
    config.gpu_options.per_process_gpu_memory_fraction = 0.5
    return tf.Session(config=config)

def parse_args(args):

    parser     = argparse.ArgumentParser(description='Evaluation script for a RetinaNet network.')
    parser.add_argument('model',              help='Path to RetinaNet model.')
    parser.add_argument('--gpu',              help='Id of the GPU to use (as reported by nvidia-smi).')
    parser.add_argument('--config',           help='Path to a configuration parameters .ini file (only used with --convert-model).')

    return parser.parse_args(args)

def load_model(model_path):

    check_keras_version()

    #if args.gpu:
    #    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    #keras.backend.tensorflow_backend.set_session(get_session())

    anchor_params = None
    backbone = 'resnet50'

    print('Loading model, this may take a second...')
    print(model_path)
    model = models.load_model(model_path, backbone_name=backbone)
    #graph = tf.compat.v1.get_default_graph()
    model = models.convert_model(model, anchor_params=anchor_params) # convert model

    # print model summary
    print(model.summary())

    return model#, graph

mask_pub = rospy.Publisher('/pyrapose/masks', Image, queue_size=10)
#def run_estimation(image, model, score_threshold, graph, frame_id):
def run_estimation(image, image_dep, model, score_threshold, threeD_boxes, model_6, model_9, model_10, model_11, model_61):
    obj_names = []
    obj_poses = []
    obj_confs = []

    image_mask = copy.deepcopy(image)
    image = preprocess_image(image)
    #image_mask = copy.deepcopy(image)

    #cv2.imwrite('/home/sthalham/retnetpose_image.jpg', image)

    if keras.backend.image_data_format() == 'channels_first':
            image = image.transpose((2, 0, 1))

    #with graph.as_default():
    boxes3D, scores, mask = model.predict_on_batch(np.expand_dims(image, axis=0))

    for inv_cls in range(scores.shape[2]):

        if inv_cls == 0:
            true_cls = 5
        elif inv_cls == 1:
            true_cls = 8
        elif inv_cls == 2:
            true_cls = 9
        elif inv_cls == 3:
            true_cls = 10
        elif inv_cls == 4:
            true_cls = 21
        #true_cat = inv_cls + 1
        #true_cls = true_cat

        cls_mask = scores[0, :, inv_cls]

        cls_indices = np.where(cls_mask > score_threshold)
        #cls_indices = np.argmax(cls_mask)
        #print(cls_mask[cls_indices])

        cls_img = image
        obj_mask = mask[0, :, inv_cls]
        if inv_cls == 0:
            obj_col = [1, 255, 255]
        elif inv_cls == 1:
            obj_col = [1, 1, 128]
        elif inv_cls == 2:
            obj_col = [255, 255, 1]
        elif inv_cls == 3:
            obj_col = [220, 245, 245]
        elif inv_cls == 4:
            obj_col = [128, 1, 1]
        cls_img = np.where(obj_mask > 0.5, 1, 0)
        cls_img = cls_img.reshape((60, 80)).astype(np.uint8)
        cls_img = np.asarray(Pilimage.fromarray(cls_img).resize((640, 480), Pilimage.NEAREST))
        depth_mask = copy.deepcopy(cls_img)
        cls_img = np.repeat(cls_img[:, :, np.newaxis], 3, 2)
        cls_img = cls_img.astype(np.uint8)
        cls_img[:, :, 0] *= obj_col[0]
        cls_img[:, :, 1] *= obj_col[1]
        cls_img[:, :, 2] *= obj_col[2]
        image_mask = np.where(cls_img > 0, cls_img, image_mask.astype(np.uint8))
        #cv2.imwrite('/stefan/mask.png', image_mask)

        #if len(cls_indices[0]) < 1:
        if len(cls_indices[0]) < 1:
            continue

        if true_cls == 5:
            name = '006_mustard_bottle'
            pcd_model = model_6
        elif true_cls == 8:
            name = '009_gelatin_box'
            pcd_model = model_9
        elif true_cls == 9:
            name = '010_potted_meat_can'
            pcd_model = model_10
        elif true_cls == 10:
            name = '011_banana'
            pcd_model = model_11
        elif true_cls == 21:
            name = '061_foam_brick'
            pcd_model = model_61
        else:
            continue

        obj_names.append(name)
        #obj_confs.append(np.sum(cls_mask[cls_indices[0]]))
        obj_confs.append(np.sum(cls_mask[cls_indices]))

        k_hyp = len(cls_indices[0])
        #k_hyp = 1
        ori_points = np.ascontiguousarray(threeD_boxes[(true_cls), :, :], dtype=np.float32)  # .reshape((8, 1, 3))
        K = np.float32([fxkin, 0., cxkin, 0., fykin, cykin, 0., 0., 1.]).reshape(3, 3)

        ##############################
        # pnp
        pose_votes = boxes3D[0, cls_indices, :]
        est_points = np.ascontiguousarray(pose_votes, dtype=np.float32).reshape((int(k_hyp * 8), 1, 2))
        obj_points = np.repeat(ori_points[np.newaxis, :, :], k_hyp, axis=0)
        obj_points = obj_points.reshape((int(k_hyp * 8), 1, 3))
        retval, orvec, otvec, inliers = cv2.solvePnPRansac(objectPoints=obj_points,
                                                            imagePoints=est_points, cameraMatrix=K,
                                                            distCoeffs=None, rvec=None, tvec=None,
                                                            useExtrinsicGuess=False, iterationsCount=300,
                                                            reprojectionError=5.0, confidence=0.99,
                                                            flags=cv2.SOLVEPNP_ITERATIVE)
        R_est, _ = cv2.Rodrigues(orvec)
        t_est = otvec[:, 0]

        if np.sum(depth_mask) > 3000 :

            print('--------------------- ICP refinement -------------------')

            print('cls: ', true_cls)
            pcd_img = np.where(depth_mask, image_dep, np.NaN)
            pcd_img = create_point_cloud(pcd_img, fxkin, fykin, cxkin, cykin, 1.0)
            pcd_img = pcd_img[~np.isnan(pcd_img).any(axis=1)]
            pcd_crop = open3d.PointCloud()
            pcd_crop.points = open3d.Vector3dVector(pcd_img)
            open3d.estimate_normals(pcd_crop, search_param=open3d.KDTreeSearchParamHybrid(radius=20.0, max_nn=30))

            guess = np.zeros((4, 4), dtype=np.float32)
            guess[:3, :3] = R_est
            guess[:3, 3] = t_est.T * 1000.0
            guess[3, :] = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32).T

            pcd_model = open3d.geometry.voxel_down_sample(pcd_model, voxel_size=5.0)
            pcd_crop = open3d.geometry.voxel_down_sample(pcd_crop, voxel_size=5.0)
            open3d.estimate_normals(pcd_crop, search_param=open3d.KDTreeSearchParamHybrid(radius=10.0, max_nn=10))
            open3d.estimate_normals(pcd_model, search_param=open3d.KDTreeSearchParamHybrid(radius=10.0, max_nn=10))

            pcd_model.transform(guess)

            # remove model vertices facing away from camera
            points_unfiltered = np.asarray(pcd_model.points)
            last_pcd_temp = []
            for i, normal in enumerate(pcd_model.normals):
                if normal[2] < 0:
                    last_pcd_temp.append(points_unfiltered[i, :])

            pcd_model.points = open3d.Vector3dVector(np.asarray(last_pcd_temp))
            open3d.estimate_normals(pcd_model, search_param=open3d.KDTreeSearchParamHybrid(
                    radius=20.0, max_nn=10))

            # align model with median depth of scene
            mean_crop = np.median(np.array(pcd_crop.points), axis=0)
            mean_model = np.mean(np.array(pcd_model.points), axis=0)
            pcd_crop_filt = copy.deepcopy(pcd_crop)

            pcd_min = guess[2, 3] - 75
            pcd_max = guess[2, 3] + 75
            new_pcd_scene = []
            for i, point in enumerate(pcd_crop.points):
                if point[2] > pcd_min and point[2] < pcd_max:
                    new_pcd_scene.append(point)

            #mean_crop = np.mean(np.array(pcd_crop_filt.points), axis=0)

            print(mean_crop, mean_crop.shape)
            #print(guess[:3, 3], guess[:3, 3].shape)
            #print(mean_crop-guess[:3, 3])
            print('euclid: ', np.linalg.norm((mean_crop-guess[:3, 3]), ord=2))
            print('num_points: ', len(new_pcd_scene))
            if len(new_pcd_scene)< 50 or np.linalg.norm((mean_crop-guess[:3, 3]), ord=2) > 75:
                print('use pcd mean')
                if len(new_pcd_scene) > 50 and np.linalg.norm((mean_crop-guess[:3, 3]), ord=2) > 75:
                    print('recalc mean')
                    pcd_crop_filt.points = open3d.Vector3dVector(np.asarray(new_pcd_scene))
                    mean_crop = np.mean(np.array(pcd_crop_filt.points), axis=0)
                pcd_diff = mean_crop - mean_model
                pcd_model.translate(pcd_diff)
                open3d.estimate_normals(pcd_model, search_param=open3d.KDTreeSearchParamHybrid(
                    radius=10.0, max_nn=10))
                guess[:3, 3] = mean_crop
            else:
                print('use pose')
                pcd_crop.points = open3d.Vector3dVector(np.asarray(new_pcd_scene))

            open3d.estimate_normals(pcd_crop, search_param=open3d.KDTreeSearchParamHybrid(
                radius=20.0, max_nn=10))

            #reg_p2p = open3d.registration.registration_icp(pcd_model, pcd_crop, 5.0, np.eye(4),
            #                                                open3d.registration.TransformationEstimationPointToPlane(), open3d.registration.ICPConvergenceCriteria(max_iteration=100))
            reg_icp = cv2.ppf_match_3d_ICP(100, tolerence=0.075, numLevels=4)
            model_points = np.asarray(pcd_model.points, dtype=np.float32)
            model_normals = np.asarray(pcd_model.normals, dtype=np.float32)
            crop_points = np.asarray(pcd_crop.points, dtype=np.float32)
            crop_normals = np.asarray(pcd_crop.normals, dtype=np.float32)
            pcd_source = np.zeros((model_points.shape[0], 6), dtype=np.float32)
            pcd_target = np.zeros((crop_points.shape[0], 6), dtype=np.float32)
            pcd_source[:, :3] = model_points * 0.001
            pcd_source[:, 3:] = model_normals
            pcd_target[:, :3] = crop_points * 0.001
            pcd_target[:, 3:] = crop_normals

            retval, residual, pose = reg_icp.registerModelToScene(pcd_source, pcd_target)

            print('residual: ', residual)
            #pcd_model.transform(reg_p2p.transformation)
            guess[:3, 3] = guess[:3, 3] * 0.001
            guess = np.matmul(pose, guess)
            R_est = guess[:3, :3]
            t_est = guess[:3, 3]

            #print('guess: ', guess)

        est_pose = np.zeros((7), dtype=np.float32)
        est_pose[:3] = t_est
        est_pose[3:] = tf3d.quaternions.mat2quat(R_est)
        obj_poses.append(est_pose)

    bridge = CvBridge()

    image_mask_msg = bridge.cv2_to_imgmsg(image_mask)

    mask_pub.publish(image_mask_msg)
    return obj_names, obj_poses, obj_confs

if __name__ == '__main__':

    # ROS params
    mesh_path = ''
    msg_topic = '/camera/rgb/image_color'
    score_threshold = 0.5
    icp_threshold = 0.15
    service_name = 'get_poses'
    try:
        model_path = rospy.get_param('/PyraPose/model_path')
    except KeyError:
        print("please set path to model! example:/home/desired/path/to/resnet_xy.h5")
    try:
        mesh_path = rospy.get_param('/PyraPose/meshes_path')
    except KeyError:
        print("please set path to meshes! example:/home/desired/path/to/meshes/")

    if rospy.has_param('/PyraPose/detection_threshold'):
        score_threshold = rospy.get_param("/PyraPose/detection_threshold")
        print('Detection threshold set to: ', score_threshold)
    if rospy.has_param('/PyraPose/image_topic'):
        msg_topic = rospy.get_param("/PyraPose/image_topic")
        print("Subscribing to msg topic: ", msg_topic)
    if rospy.has_param('/PyraPose/icp_threshold'):
        icp_threshold = rospy.get_param("/PyraPose/icp_threshold")
        print("icp threshold set to: ", icp_threshold)
    if rospy.has_param('/PyraPose/service_call'):
        service_name = rospy.get_param("/PyraPose/service_call")
        print("service call set to: ", service_name)

    #model, graph = load_model(model_path)
    model = load_model(model_path)
    try:
        if rospy.get_param('/PyraPose/node_type') == 'continuous':
            print("node type set to continuous")
            pose_estimation = PoseEstimationClass(model, mesh_path, score_threshold, msg_topic)#, graph)
        elif rospy.get_param('/PyraPose/node_type') == 'service':
            print("node type set to service")
            pose_estimation = PoseEstimationServer(model, mesh_path, score_threshold, msg_topic, service_name)
    except KeyError:
        print("node_type should either be continuous or service.")
    rospy.init_node('PyraPose', anonymous=True)

    rospy.spin()

ext: py

max_stars_repo_name: sThalham/UDAPose

max_issues_repo_name: sThalham/UDAPose

max_forks_repo_name: sThalham/UDAPose

reason: import numpy,from scipy


text:
# View more 3_python 2_tensorflow_old on my Youtube and Youku channel!!!

# Youtube video tutorial: https://www.youtube.com/channel/UCdyjiB5H8Pu7aDTNVXTTpcg
# Youku video tutorial: http://i.youku.com/pythontutorial

# 3 - simple plot
"""
Please note, this script is for python3+.
If you are using python2+, please modify it accordingly.
"""

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(-1, 1, 50)
y = 2*x + 1
# y = x**2
plt.plot(x, y)
plt.show()

reason: import numpy


text:
# Copyright 2020 The MuLT Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import pandas as pd
import numpy as np

def classification_metrics(tn, fp, fn, tp):

    sensitivity = (tp / float(tp + fn)) if tp + fn > 0 else 1

    precision = (tp / float(tp + fp)) if tp + fp > 0 else 1

    specificity = (tn / float(tn + fp)) if tn + fp > 0 else 1

    accuracy = (tp + tn) / (tp + tn + fp + fn)

    return {'accuracy': accuracy, 'precision': precision, 'sensitivity': sensitivity, 'specificity': specificity}

def ks_score(y_true, y_hat):

    try:

        deciles = [-1] + list(np.quantile(y_hat, [.1, .2, .3, .4, .5, .6, .7, .8, .9, 1]))

        result = []

        for i in y_hat:
            for index, (a, b) in enumerate(zip(deciles[:-1], deciles[1:])):
                if a < i <= b:
                    result.append(index + 1)
                    continue

        result = pd.DataFrame({'decile': result})

        result['positive'] = y_true
        result['negative'] = 1 - y_true

        negative_sum = result['negative'].sum()
        positive_sum = result['positive'].sum()

        result = result.groupby(by=['decile']).sum()

        result['positive_percentage'] = result['positive'] / positive_sum
        result['negative_percentage'] = result['negative'] / negative_sum

        result['positive_percentage_cumsum'] = result['positive_percentage'].cumsum()
        result['negative_percentage_cumsum'] = result['negative_percentage'].cumsum()

        result['ks'] = np.abs(result['positive_percentage_cumsum'] - result['negative_percentage_cumsum'])

        return result['ks'].max()

    except ValueError:
        return 0.0

max_stars_repo_name: urielcaire/mult

max_issues_repo_name: urielcaire/mult

max_forks_repo_name: urielcaire/mult


text:
import cv2
import numpy as np;

def analyzeImage(imageNumber):
    #setup
    colonyCount = 0
#     fileName = "/home/pi/Documents/deltaImageProcessor/images/"\
#             + str(imageNumber) + ".jpg"

    fileName = "/home/pi/Documents/BlobTest.jpg"
    image = cv2.imread(fileName, cv2.IMREAD_GRAYSCALE)  #import jpeg

    #Setup SimpleBlobDetector parameters.
    params = cv2.SimpleBlobDetector_Params()

    # Change thresholds
    params.minThreshold = 10
    params.maxThreshold = 200

    # Filter by Area.
    params.filterByArea = True
    params.minArea = 1500

    # Filter by Circularity
    params.filterByCircularity = True
    params.minCircularity = 0.1
    params.maxCircularity = .785

    # Filter by Convexity
    params.filterByConvexity = True
    params.minConvexity = 0.87

    # Filter by Inertia
    params.filterByInertia = True
    params.minInertiaRatio = 0.05

    # Create a detector with the parameters
    detector = cv2.SimpleBlobDetector_create(params)

    # Detect blobs.
    keypoints = detector.detect(image)
    colonyCount = len(keypoints)
    # Draw detected blobs as red circles.
    # cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS ensures the size of the circle corresponds to the size of blob
    im_with_keypoints = cv2.drawKeypoints(image, keypoints, np.array([]), (0,0,255), cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)

    # Show keypoints
    cv2.imshow("Keypoints", im_with_keypoints)
    cv2.waitKey(0)
    return colonyCount

max_stars_repo_name: zbendt/ECE-Capstone-Project

max_issues_repo_name: zbendt/ECE-Capstone-Project

max_forks_repo_name: zbendt/ECE-Capstone-Project


text:
import unittest
import numpy as np
from numpy.testing import assert_array_equal

from graphs.base.pairs import EdgePairGraph, SymmEdgePairGraph

PAIRS = np.array([[0,1],[0,2],[1,1],[2,1],[3,3]])
ADJ = [[0,1,1,0],
       [0,1,0,0],
       [0,1,0,0],
       [0,0,0,1]]

class TestEdgePairGraph(unittest.TestCase):
  def setUp(self):
    self.epg = EdgePairGraph(PAIRS)

  def test_pairs(self):
    self.assert_(self.epg.pairs(copy=False) is PAIRS)
    P = self.epg.pairs(copy=True)
    self.assert_(P is not PAIRS)
    assert_array_equal(P, PAIRS)
    # test the directed case
    P = self.epg.pairs(directed=False)
    assert_array_equal(P, [[0,1],[0,2],[1,1],[1,2],[3,3]])

  def test_matrix(self):
    M = self.epg.matrix()
    assert_array_equal(M.toarray(), ADJ)
    M = self.epg.matrix('dense')
    assert_array_equal(M, ADJ)
    M = self.epg.matrix('csr')
    self.assertEqual(M.format, 'csr')
    assert_array_equal(M.toarray(), ADJ)

  def test_self_edges(self):
    self.epg.add_self_edges()
    expected = self.epg.pairs()
    # Ensure that calling it again does the right thing.
    self.epg.add_self_edges()
    assert_array_equal(self.epg.pairs(), expected)

  def test_symmetrize(self):
    # Check that copy=True doesn't change anything
    self.epg.symmetrize(copy=True)
    assert_array_equal(self.epg.matrix('dense'), ADJ)

class TestSymmEdgePairGraph(unittest.TestCase):
  def setUp(self):
    self.G = SymmEdgePairGraph(PAIRS)

  def test_copy(self):
    gg = self.G.copy()
    self.assertIsNot(gg, self.G)
    assert_array_equal(gg.matrix('dense'), self.G.matrix('dense'))
    assert_array_equal(gg.pairs(), self.G.pairs())

  def test_pairs(self):
    expected = [[0,1], [0,2], [1,0], [1,1], [1,2], [2,0], [2,1], [3,3]]
    P = self.G.pairs()
    assert_array_equal(sorted(P.tolist()), expected)
    # test the directed case
    P = self.G.pairs(directed=False)
    assert_array_equal(P, [[0,1],[0,2],[1,1],[1,2],[3,3]])

  def test_symmetrize(self):
    self.assertIs(self.G.symmetrize(copy=False), self.G)
    S = self.G.symmetrize(copy=True)
    self.assertIsNot(S, self.G)
    assert_array_equal(S.matrix('dense'), self.G.matrix('dense'))

if __name__ == '__main__':
  unittest.main()

max_stars_repo_name: vishalbelsare/graphs

max_issues_repo_name: perimosocordiae/graphs

max_forks_repo_name: perimosocordiae/graphs

reason: import numpy,from numpy