sonusai 0.16.1__py3-none-any.whl → 0.17.0__py3-none-any.whl

sonusai/audiofe.py

@@ -12,7 +12,7 @@ options:
     -m MODEL, --model MODEL         PL model .py file path.
     -k CKPT, --checkpoint CKPT      PL checkpoint file with weights.
     -a ASR, --asr ASR               ASR method to use.
-    -w WMODEL, --whisper WMODEL     Whisper model used in aixplain_whisper and whisper methods. [default: tiny].
+    -w WMODEL, --whisper WMODEL     Model used in whisper, aixplain_whisper and faster_whisper methods. [default: tiny].
 
 Aaware SonusAI Audio Front End.
 
@@ -29,7 +29,7 @@ audiofe_capture_<TIMESTAMP>.png and predict data (time-domain signal and feature
 audiofe_predict_<TIMESTAMP>.png.
 
 If an ASR is specified, run ASR on the captured audio and print the results. In addition, if a model was also specified,
-run ASR on the predict audio and print the results.
+run ASR on the predict audio and print the results.  Examples: faster_whisper, google,
 
 If the debug option is enabled, write capture audio, feature, reconstruct audio, predict, and predict audio to
 audiofe_<TIMESTAMP>.h5.
@@ -79,6 +79,7 @@ def main() -> None:
     import torch
     from docopt import printable_usage
     from sonusai_torchl.utils import load_torchl_ckpt_model
+    from sonusai.utils.onnx_utils import load_ort_session
 
     from sonusai import create_file_handler
     from sonusai import initial_log_messages
@@ -102,9 +103,32 @@ def main() -> None:
     predict_png = predict_name + '.png'
     h5_name = f'audiofe_{ts}.h5'
 
-    if model_name is not None and ckpt_name is None:
-        print(printable_usage(trim_docstring(__doc__)))
-        exit(1)
+    if model_name is not None:
+        from os.path import splitext
+        if splitext(model_name)[1] == '.onnx':
+            session, options, model_root, hparams, sess_inputs, sess_outputs = load_ort_session(model_name)
+            if hparams is None:
+                logger.error(f'Error: onnx model does not have required SonusAI hyper-parameters, can not proceed.')
+                raise SystemExit(1)
+            feature_mode = hparams["feature"]
+            model_is_onnx = True
+            in0name = sess_inputs[0].name
+            in0type = sess_inputs[0].type
+            out0name = sess_outputs[0].name
+            out_names = [n.name for n in session.get_outputs()]
+            if in0type.find('float16') != -1:
+                model_is_fp16 = True
+                logger.info(f'Detected input of float16, converting all feature inputs to that type.')
+            else:
+                model_is_fp16 = False
+        else:
+            model_is_onnx = False
+            if ckpt_name is None:
+                print(printable_usage(trim_docstring(__doc__)))
+                exit(1)
+            model = load_torchl_ckpt_model(model_name=model_name, ckpt_name=ckpt_name)
+            feature_mode = model.hparams.feature
+            model.eval()
 
     # Setup logging file
     create_file_handler('audiofe.log')
@@ -129,26 +153,25 @@ def main() -> None:
         except ValueError as e:
             logger.exception(e)
             return
+        # Only write if capture, not for file input
+        write_wav(capture_wav, capture_audio, SAMPLE_RATE)
+        logger.info('')
+        logger.info(f'Wrote capture audio with shape {capture_audio.shape} to {capture_wav}')
 
-    write_wav(capture_wav, capture_audio, SAMPLE_RATE)
-    logger.info('')
-    logger.info(f'Wrote capture audio with shape {capture_audio.shape} to {capture_wav}')
     if debug:
         with h5py.File(h5_name, 'a') as f:
             if 'capture_audio' in f:
                 del f['capture_audio']
             f.create_dataset('capture_audio', data=capture_audio)
-        logger.info(f'Wrote capture audio with shape {capture_audio.shape} to {h5_name}')
+        logger.info(f'Wrote capture feature data with shape {capture_audio.shape} to {h5_name}')
 
     if asr_name is not None:
+        logger.info(f'Running ASR on captured audio with {asr_name} ...')
         capture_asr = calc_asr(capture_audio, engine=asr_name, whisper_model_name=whisper_name).text
         logger.info(f'Capture audio ASR: {capture_asr}')
 
     if model_name is not None:
-        model = load_torchl_ckpt_model(model_name=model_name, ckpt_name=ckpt_name)
-        model.eval()
-
-        feature = get_feature_from_audio(audio=capture_audio, feature_mode=model.hparams.feature)
+        feature = get_feature_from_audio(audio=capture_audio, feature_mode=feature_mode)  #frames x stride x feat_params
         save_figure(capture_png, capture_audio, feature)
         logger.info(f'Wrote capture plots to {capture_png}')
 
@@ -159,9 +182,20 @@ def main() -> None:
                 f.create_dataset('feature', data=feature)
             logger.info(f'Wrote feature with shape {feature.shape} to {h5_name}')
 
-        with torch.no_grad():
-            # model wants batch x timesteps x feature_parameters
-            predict = model(torch.tensor(feature).permute((1, 0, 2))).permute(1, 0, 2).numpy()
+        if model_is_onnx:
+            # run ort session, wants i.e. batch x tsteps x feat_params, outputs numpy BxTxFP or BxFP
+            # Note full reshape not needed here since we assume speech enhanement type model, so a transpose suffices
+            if model_is_fp16:
+                feature = np.float16(feature)
+            # run inference, ort session wants i.e. batch x tsteps x feat_params, outputs numpy BxTxFP or BxFP
+            predict = np.transpose(session.run(out_names, {in0name: np.transpose(feature,(1,0,2))})[0],(1,0,2))
+        else:
+            with torch.no_grad():
+                # model wants batch x timesteps x feature_parameters
+                predict = model(torch.tensor(feature).permute((1, 0, 2))).permute(1, 0, 2).numpy()
+
+
+
         if debug:
             with h5py.File(h5_name, 'a') as f:
                 if 'predict' in f:
@@ -169,7 +203,7 @@ def main() -> None:
                 f.create_dataset('predict', data=predict)
             logger.info(f'Wrote predict with shape {predict.shape} to {h5_name}')
 
-        predict_audio = get_audio_from_feature(feature=predict, feature_mode=model.hparams.feature)
+        predict_audio = get_audio_from_feature(feature=predict, feature_mode=feature_mode)
         write_wav(predict_wav, predict_audio, SAMPLE_RATE)
         logger.info(f'Wrote predict audio with shape {predict_audio.shape} to {predict_wav}')
         if debug:
@@ -183,6 +217,7 @@ def main() -> None:
         logger.info(f'Wrote predict plots to {predict_png}')
 
         if asr_name is not None:
+            logger.info(f'Running ASR on model-enhanced audio with {asr_name} ...')
             predict_asr = calc_asr(predict_audio, engine=asr_name, whisper_model_name=whisper_name).text
             logger.info(f'Predict audio ASR: {predict_asr}')