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

sonusai/audiofe.py

@@ -1,18 +1,17 @@
 """sonusai audiofe
 
-usage: audiofe [-hvds] [--version] [-i INPUT] [-l LENGTH] [-m MODEL] [-k CKPT] [-a ASR] [-w WMODEL]
+usage: audiofe [-hvds] [--version] [-i INPUT] [-l LENGTH] [-m MODEL] [-a ASR] [-w WMODEL]
 
 options:
     -h, --help
     -v, --verbose                   Be verbose.
     -d, --debug                     Write debug data to H5 file.
-    -s, --show                      Show a list of available audio inputs.
+    -s, --show                      Display a list of available audio inputs.
     -i INPUT, --input INPUT         Input audio.
     -l LENGTH, --length LENGTH      Length of audio in seconds. [default: -1].
-    -m MODEL, --model MODEL         PL model .py file path.
-    -k CKPT, --checkpoint CKPT      PL checkpoint file with weights.
+    -m MODEL, --model MODEL         ONNX model.
     -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 +28,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.
@@ -66,7 +65,6 @@ def main() -> None:
     length = float(args['--length'])
     input_name = args['--input']
     model_name = args['--model']
-    ckpt_name = args['--checkpoint']
     asr_name = args['--asr']
     whisper_name = args['--whisper']
     debug = args['--debug']
@@ -76,9 +74,6 @@ def main() -> None:
 
     import h5py
     import pyaudio
-    import torch
-    from docopt import printable_usage
-    from sonusai_torchl.utils import load_torchl_ckpt_model
 
     from sonusai import create_file_handler
     from sonusai import initial_log_messages
@@ -90,7 +85,7 @@ def main() -> None:
     from sonusai.utils import calc_asr
     from sonusai.utils import create_timestamp
     from sonusai.utils import get_input_devices
-    from sonusai.utils import trim_docstring
+    from sonusai.utils import load_ort_session
     from sonusai.utils import write_wav
 
     ts = create_timestamp()
@@ -102,10 +97,6 @@ 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)
-
     # Setup logging file
     create_file_handler('audiofe.log')
     update_console_handler(verbose)
@@ -129,26 +120,35 @@ def main() -> None:
         except ValueError as e:
             logger.exception(e)
             return
+        # Only write if capture from device, 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)
+        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 hyperparameters, cannot proceed.')
+            raise SystemExit(1)
+        feature_mode = hparams.feature
+        in0name = sess_inputs[0].name
+        in0type = sess_inputs[0].type
+        out_names = [n.name for n in session.get_outputs()]
+
+        # frames x stride x feat_params
+        feature = get_feature_from_audio(audio=capture_audio, feature_mode=feature_mode)
         save_figure(capture_png, capture_audio, feature)
         logger.info(f'Wrote capture plots to {capture_png}')
 
@@ -159,9 +159,14 @@ 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 in0type.find('float16') != -1:
+            logger.info(f'Detected input of float16, converting all feature inputs to that type.')
+            feature = np.float16(feature)  # type: ignore
+
+        # Run inference, ort session wants batch x timesteps x feat_params, outputs numpy BxTxFP or BxFP
+        # Note full reshape not needed here since we assume speech enhancement type model, so a transpose suffices
+        predict = np.transpose(session.run(out_names, {in0name: np.transpose(feature, (1, 0, 2))})[0], (1, 0, 2))
+
         if debug:
             with h5py.File(h5_name, 'a') as f:
                 if 'predict' in f:
@@ -169,7 +174,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 +188,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}')