| | |
|---|
| | | # -*- encoding: utf-8 -*- |
|---|
| | | from pathlib import Path |
|---|
| | | from typing import Any, Dict, Iterable, List, NamedTuple, Set, Tuple, Union |
|---|
| | | import copy |
|---|
| | | |
|---|
| | | import numpy as np |
|---|
| | | from typeguard import check_argument_types |
|---|
| | |
|---|
| | | cmvn = np.array([means, vars]) |
|---|
| | | return cmvn |
|---|
| | | |
|---|
| | | |
|---|
| | | class WavFrontendOnline(WavFrontend): |
|---|
| | | def __init__(self, **kwargs): |
|---|
| | | super().__init__(**kwargs) |
|---|
| | | # self.fbank_fn = knf.OnlineFbank(self.opts) |
|---|
| | | # add variables |
|---|
| | | self.frame_sample_length = int(self.opts.frame_opts.frame_length_ms * self.opts.frame_opts.samp_freq / 1000) |
|---|
| | | self.frame_shift_sample_length = int(self.opts.frame_opts.frame_shift_ms * self.opts.frame_opts.samp_freq / 1000) |
|---|
| | | self.waveform = None |
|---|
| | | self.reserve_waveforms = None |
|---|
| | | self.input_cache = None |
|---|
| | | self.lfr_splice_cache = [] |
|---|
| | | |
|---|
| | | @staticmethod |
|---|
| | | # inputs has catted the cache |
|---|
| | | def apply_lfr(inputs: np.ndarray, lfr_m: int, lfr_n: int, is_final: bool = False) -> Tuple[ |
|---|
| | | np.ndarray, np.ndarray, int]: |
|---|
| | | """ |
|---|
| | | Apply lfr with data |
|---|
| | | """ |
|---|
| | | |
|---|
| | | LFR_inputs = [] |
|---|
| | | T = inputs.shape[0] # include the right context |
|---|
| | | T_lfr = int(np.ceil((T - (lfr_m - 1) // 2) / lfr_n)) # minus the right context: (lfr_m - 1) // 2 |
|---|
| | | splice_idx = T_lfr |
|---|
| | | for i in range(T_lfr): |
|---|
| | | if lfr_m <= T - i * lfr_n: |
|---|
| | | LFR_inputs.append((inputs[i * lfr_n:i * lfr_n + lfr_m]).reshape(1, -1)) |
|---|
| | | else: # process last LFR frame |
|---|
| | | if is_final: |
|---|
| | | num_padding = lfr_m - (T - i * lfr_n) |
|---|
| | | frame = (inputs[i * lfr_n:]).reshape(-1) |
|---|
| | | for _ in range(num_padding): |
|---|
| | | frame = np.hstack((frame, inputs[-1])) |
|---|
| | | LFR_inputs.append(frame) |
|---|
| | | else: |
|---|
| | | # update splice_idx and break the circle |
|---|
| | | splice_idx = i |
|---|
| | | break |
|---|
| | | splice_idx = min(T - 1, splice_idx * lfr_n) |
|---|
| | | lfr_splice_cache = inputs[splice_idx:, :] |
|---|
| | | LFR_outputs = np.vstack(LFR_inputs) |
|---|
| | | return LFR_outputs.astype(np.float32), lfr_splice_cache, splice_idx |
|---|
| | | |
|---|
| | | @staticmethod |
|---|
| | | def compute_frame_num(sample_length: int, frame_sample_length: int, frame_shift_sample_length: int) -> int: |
|---|
| | | frame_num = int((sample_length - frame_sample_length) / frame_shift_sample_length + 1) |
|---|
| | | return frame_num if frame_num >= 1 and sample_length >= frame_sample_length else 0 |
|---|
| | | |
|---|
| | | |
|---|
| | | def fbank( |
|---|
| | | self, |
|---|
| | | input: np.ndarray, |
|---|
| | | input_lengths: np.ndarray |
|---|
| | | ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: |
|---|
| | | self.fbank_fn = knf.OnlineFbank(self.opts) |
|---|
| | | batch_size = input.shape[0] |
|---|
| | | if self.input_cache is None: |
|---|
| | | self.input_cache = np.empty((batch_size, 0), dtype=np.float32) |
|---|
| | | input = np.concatenate((self.input_cache, input), axis=1) |
|---|
| | | frame_num = self.compute_frame_num(input.shape[-1], self.frame_sample_length, self.frame_shift_sample_length) |
|---|
| | | # update self.in_cache |
|---|
| | | self.input_cache = input[:, -(input.shape[-1] - frame_num * self.frame_shift_sample_length):] |
|---|
| | | waveforms = np.empty(0, dtype=np.int16) |
|---|
| | | feats_pad = np.empty(0, dtype=np.float32) |
|---|
| | | feats_lens = np.empty(0, dtype=np.int32) |
|---|
| | | if frame_num: |
|---|
| | | waveforms = [] |
|---|
| | | feats = [] |
|---|
| | | feats_lens = [] |
|---|
| | | for i in range(batch_size): |
|---|
| | | waveform = input[i] |
|---|
| | | waveforms.append( |
|---|
| | | waveform[:((frame_num - 1) * self.frame_shift_sample_length + self.frame_sample_length)]) |
|---|
| | | waveform = waveform * (1 << 15) |
|---|
| | | |
|---|
| | | self.fbank_fn.accept_waveform(self.opts.frame_opts.samp_freq, waveform.tolist()) |
|---|
| | | frames = self.fbank_fn.num_frames_ready |
|---|
| | | mat = np.empty([frames, self.opts.mel_opts.num_bins]) |
|---|
| | | for i in range(frames): |
|---|
| | | mat[i, :] = self.fbank_fn.get_frame(i) |
|---|
| | | feat = mat.astype(np.float32) |
|---|
| | | feat_len = np.array(mat.shape[0]).astype(np.int32) |
|---|
| | | feats.append(mat) |
|---|
| | | feats_lens.append(feat_len) |
|---|
| | | |
|---|
| | | waveforms = np.stack(waveforms) |
|---|
| | | feats_lens = np.array(feats_lens) |
|---|
| | | feats_pad = np.array(feats) |
|---|
| | | self.fbanks = feats_pad |
|---|
| | | self.fbanks_lens = copy.deepcopy(feats_lens) |
|---|
| | | return waveforms, feats_pad, feats_lens |
|---|
| | | |
|---|
| | | def get_fbank(self) -> Tuple[np.ndarray, np.ndarray]: |
|---|
| | | return self.fbanks, self.fbanks_lens |
|---|
| | | |
|---|
| | | def lfr_cmvn( |
|---|
| | | self, |
|---|
| | | input: np.ndarray, |
|---|
| | | input_lengths: np.ndarray, |
|---|
| | | is_final: bool = False |
|---|
| | | ) -> Tuple[np.ndarray, np.ndarray, List[int]]: |
|---|
| | | batch_size = input.shape[0] |
|---|
| | | feats = [] |
|---|
| | | feats_lens = [] |
|---|
| | | lfr_splice_frame_idxs = [] |
|---|
| | | for i in range(batch_size): |
|---|
| | | mat = input[i, :input_lengths[i], :] |
|---|
| | | lfr_splice_frame_idx = -1 |
|---|
| | | if self.lfr_m != 1 or self.lfr_n != 1: |
|---|
| | | # update self.lfr_splice_cache in self.apply_lfr |
|---|
| | | mat, self.lfr_splice_cache[i], lfr_splice_frame_idx = self.apply_lfr(mat, self.lfr_m, self.lfr_n, |
|---|
| | | is_final) |
|---|
| | | if self.cmvn_file is not None: |
|---|
| | | mat = self.apply_cmvn(mat) |
|---|
| | | feat_length = mat.shape[0] |
|---|
| | | feats.append(mat) |
|---|
| | | feats_lens.append(feat_length) |
|---|
| | | lfr_splice_frame_idxs.append(lfr_splice_frame_idx) |
|---|
| | | |
|---|
| | | feats_lens = np.array(feats_lens) |
|---|
| | | feats_pad = np.array(feats) |
|---|
| | | return feats_pad, feats_lens, lfr_splice_frame_idxs |
|---|
| | | |
|---|
| | | |
|---|
| | | def extract_fbank( |
|---|
| | | self, input: np.ndarray, input_lengths: np.ndarray, is_final: bool = False |
|---|
| | | ) -> Tuple[np.ndarray, np.ndarray]: |
|---|
| | | batch_size = input.shape[0] |
|---|
| | | assert batch_size == 1, 'we support to extract feature online only when the batch size is equal to 1 now' |
|---|
| | | waveforms, feats, feats_lengths = self.fbank(input, input_lengths) # input shape: B T D |
|---|
| | | if feats.shape[0]: |
|---|
| | | self.waveforms = waveforms if self.reserve_waveforms is None else np.concatenate( |
|---|
| | | (self.reserve_waveforms, waveforms), axis=1) |
|---|
| | | if not self.lfr_splice_cache: |
|---|
| | | for i in range(batch_size): |
|---|
| | | self.lfr_splice_cache.append(np.expand_dims(feats[i][0, :], axis=0).repeat((self.lfr_m - 1) // 2, axis=0)) |
|---|
| | | |
|---|
| | | if feats_lengths[0] + self.lfr_splice_cache[0].shape[0] >= self.lfr_m: |
|---|
| | | lfr_splice_cache_np = np.stack(self.lfr_splice_cache) # B T D |
|---|
| | | feats = np.concatenate((lfr_splice_cache_np, feats), axis=1) |
|---|
| | | feats_lengths += lfr_splice_cache_np[0].shape[0] |
|---|
| | | frame_from_waveforms = int( |
|---|
| | | (self.waveforms.shape[1] - self.frame_sample_length) / self.frame_shift_sample_length + 1) |
|---|
| | | minus_frame = (self.lfr_m - 1) // 2 if self.reserve_waveforms is None else 0 |
|---|
| | | feats, feats_lengths, lfr_splice_frame_idxs = self.lfr_cmvn(feats, feats_lengths, is_final) |
|---|
| | | if self.lfr_m == 1: |
|---|
| | | self.reserve_waveforms = None |
|---|
| | | else: |
|---|
| | | reserve_frame_idx = lfr_splice_frame_idxs[0] - minus_frame |
|---|
| | | # print('reserve_frame_idx: ' + str(reserve_frame_idx)) |
|---|
| | | # print('frame_frame: ' + str(frame_from_waveforms)) |
|---|
| | | self.reserve_waveforms = self.waveforms[:, reserve_frame_idx * self.frame_shift_sample_length:frame_from_waveforms * self.frame_shift_sample_length] |
|---|
| | | sample_length = (frame_from_waveforms - 1) * self.frame_shift_sample_length + self.frame_sample_length |
|---|
| | | self.waveforms = self.waveforms[:, :sample_length] |
|---|
| | | else: |
|---|
| | | # update self.reserve_waveforms and self.lfr_splice_cache |
|---|
| | | self.reserve_waveforms = self.waveforms[:, |
|---|
| | | :-(self.frame_sample_length - self.frame_shift_sample_length)] |
|---|
| | | for i in range(batch_size): |
|---|
| | | self.lfr_splice_cache[i] = np.concatenate((self.lfr_splice_cache[i], feats[i]), axis=0) |
|---|
| | | return np.empty(0, dtype=np.float32), feats_lengths |
|---|
| | | else: |
|---|
| | | if is_final: |
|---|
| | | self.waveforms = waveforms if self.reserve_waveforms is None else self.reserve_waveforms |
|---|
| | | feats = np.stack(self.lfr_splice_cache) |
|---|
| | | feats_lengths = np.zeros(batch_size, dtype=np.int32) + feats.shape[1] |
|---|
| | | feats, feats_lengths, _ = self.lfr_cmvn(feats, feats_lengths, is_final) |
|---|
| | | if is_final: |
|---|
| | | self.cache_reset() |
|---|
| | | return feats, feats_lengths |
|---|
| | | |
|---|
| | | def get_waveforms(self): |
|---|
| | | return self.waveforms |
|---|
| | | |
|---|
| | | def cache_reset(self): |
|---|
| | | self.fbank_fn = knf.OnlineFbank(self.opts) |
|---|
| | | self.reserve_waveforms = None |
|---|
| | | self.input_cache = None |
|---|
| | | self.lfr_splice_cache = [] |
|---|
| | | |
|---|
| | | def load_bytes(input): |
|---|
| | | middle_data = np.frombuffer(input, dtype=np.int16) |
|---|
| | | middle_data = np.asarray(middle_data) |
|---|
| | |
|---|
| | | return feat, feat_len |
|---|
| | | |
|---|
| | | if __name__ == '__main__': |
|---|
| | | test() |
|---|
| | | test() |
|---|
