sonusai 0.18.9__py3-none-any.whl → 0.19.6__py3-none-any.whl

sonusai/__init__.py

@@ -2,11 +2,10 @@ import logging
 from importlib import metadata
 from os.path import dirname
 
-from pyaaware import TorchForwardTransform
-from pyaaware import TorchInverseTransform
+from rich.logging import RichHandler
+from rich.traceback import install
 
-ForwardTransform = TorchForwardTransform
-InverseTransform = TorchInverseTransform
+install(show_locals=True)
 
 __version__ = metadata.version(__package__)
 BASEDIR = dirname(__file__)
@@ -19,34 +18,26 @@ commands_doc = """
    genmetrics                   Generate mixture metrics data
    genmix                       Generate mixture and truth data
    genmixdb                     Generate a mixture database
-   gentcst                      Generate target configuration from a subdirectory tree
    lsdb                         List information about a mixture database
    mkwav                        Make WAV files from a mixture database
    onnx_predict                 Run ONNX predict on a trained model
-   plot                         Plot mixture data
    summarize_metric_spenh       Summarize speech enhancement and analysis results
-   tplot                        Plot truth data
    vars                         List custom SonusAI variables
 """
 
 # create logger
-logger = logging.getLogger('sonusai')
+logger = logging.getLogger("sonusai")
 logger.setLevel(logging.DEBUG)
-formatter = logging.Formatter('%(message)s')
-console_handler = logging.StreamHandler()
+formatter = logging.Formatter("%(message)s")
+console_handler = RichHandler(show_level=False, show_path=False, show_time=False)
 console_handler.setLevel(logging.DEBUG)
 console_handler.setFormatter(formatter)
 logger.addHandler(console_handler)
 
 
-class SonusAIError(Exception):
-    def __init__(self, value):
-        logger.error(value)
-
-
 # create file handler
 def create_file_handler(filename: str) -> None:
-    fh = logging.FileHandler(filename=filename, mode='w')
+    fh = logging.FileHandler(filename=filename, mode="w")
     fh.setLevel(logging.DEBUG)
     fh.setFormatter(formatter)
     logger.addHandler(fh)
@@ -61,7 +52,7 @@ def update_console_handler(verbose: bool) -> None:
 
 
 # write initial log message
-def initial_log_messages(name: str, subprocess: str = None) -> None:
+def initial_log_messages(name: str, subprocess: str | None = None) -> None:
     from datetime import datetime
     from getpass import getuser
     from os import getcwd
@@ -69,24 +60,24 @@ def initial_log_messages(name: str, subprocess: str = None) -> None:
     from sys import argv
 
     if subprocess is None:
-        logger.info(f'SonusAI {__version__}')
+        logger.info(f"SonusAI {__version__}")
     else:
-        logger.info(f'SonusAI {subprocess}')
-    logger.info(f'{name}')
-    logger.info('')
-    logger.debug(f'Host:      {gethostname()}')
-    logger.debug(f'User:      {getuser()}')
-    logger.debug(f'Directory: {getcwd()}')
-    logger.debug(f'Date:      {datetime.now()}')
-    logger.debug(f'Command:   {" ".join(argv)}')
-    logger.debug('')
+        logger.info(f"SonusAI {subprocess}")
+    logger.info(f"{name}")
+    logger.info("")
+    logger.debug(f"Host:      {gethostname()}")
+    logger.debug(f"User:      {getuser()}")
+    logger.debug(f"Directory: {getcwd()}")
+    logger.debug(f"Date:      {datetime.now()}")
+    logger.debug(f"Command:   {' '.join(argv)}")
+    logger.debug("")
 
 
 def commands_list(doc: str = commands_doc) -> list[str]:
-    lines = doc.split('\n')
+    lines = doc.split("\n")
     commands = []
     for line in lines:
-        command = line.strip().split(' ').pop(0)
+        command = line.strip().split(" ").pop(0)
         if command:
             commands.append(command)
     return commands

sonusai/aawscd_probwrite.py

@@ -12,19 +12,20 @@ aawscd_probwrite connects to an Aaware platform running aawscd and writes the so
 probability output to an HDF5 file.
 
 """
+
 from threading import Condition
-from typing import Optional
 
 import numpy as np
-from tqdm import tqdm
 
-CLIENT: str = 'aawscd_probwrite'
-TOPIC: str = 'aawscd/sc/prob'
+from sonusai.utils import track
+
+CLIENT: str = "aawscd_probwrite"
+TOPIC: str = "aawscd/sc/prob"
 DONE: Condition = Condition()
 FRAMES: int = 10
 FRAME_COUNT: int = 0
-DATA: Optional[np.ndarray] = None
-PROGRESS: Optional[tqdm] = None
+DATA: np.ndarray | None = None
+PROGRESS: track | None = None
 
 
 def shutdown(_signum, _frame) -> None:
@@ -75,8 +76,8 @@ def on_message(_client, _userdata, message):
 
     global TOPIC
     if mqtt.topic_matches_sub(TOPIC, message.topic):
-        payload = yaml.safe_load(str(message.payload.decode('utf-8')))
-        prob = parse_prob(payload['prob'])
+        payload = yaml.safe_load(str(message.payload.decode("utf-8")))
+        prob = parse_prob(payload["prob"])
 
         global DATA
         global FRAMES
@@ -99,20 +100,19 @@ def on_message(_client, _userdata, message):
 def main() -> None:
     from docopt import docopt
 
-    args = docopt(__doc__, version='1.0.0', options_first=True)
+    args = docopt(__doc__, version="1.0.0", options_first=True)
 
     import signal
 
     import h5py
     import paho.mqtt.client as mqtt
-    from tqdm import tqdm
 
-    machine = args['--machine']
+    machine = args["--machine"]
 
     global FRAMES
-    FRAMES = int(args['--frames'])
+    FRAMES = int(args["--frames"])
 
-    file = args['FILE']
+    file = args["FILE"]
 
     signal.signal(signal.SIGINT, shutdown)
     signal.signal(signal.SIGTERM, shutdown)
@@ -126,7 +126,7 @@ def main() -> None:
     client.subscribe(topic=TOPIC)
 
     global PROGRESS
-    PROGRESS = tqdm(total=FRAMES, desc=file)
+    PROGRESS = track(total=FRAMES, desc=file)
 
     with DONE:
         DONE.wait()
@@ -138,11 +138,11 @@ def main() -> None:
     client.disconnect()
 
     global DATA
-    with h5py.File(file, 'w') as f:
-        f.create_dataset(name='prob', data=DATA)
+    with h5py.File(file, "w") as f:
+        f.create_dataset(name="prob", data=DATA)
 
-    print(f'Wrote {file}')
+    print(f"Wrote {file}")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

sonusai/audiofe.py

@@ -34,6 +34,7 @@ If the debug option is enabled, write capture audio, feature, reconstruct audio,
 audiofe_<TIMESTAMP>.h5.
 
 """
+
 import signal
 
 import numpy as np
@@ -46,7 +47,7 @@ def signal_handler(_sig, _frame):
 
     from sonusai import logger
 
-    logger.info('Canceled due to keyboard interrupt')
+    logger.info("Canceled due to keyboard interrupt")
     sys.exit(1)
 
 
@@ -61,14 +62,14 @@ def main() -> None:
 
     args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
 
-    verbose = args['--verbose']
-    length = float(args['--length'])
-    input_name = args['--input']
-    model_name = args['--model']
-    asr_name = args['--asr']
-    whisper_name = args['--whisper']
-    debug = args['--debug']
-    show = args['--show']
+    verbose = args["--verbose"]
+    length = float(args["--length"])
+    input_name = args["--input"]
+    model_name = args["--model"]
+    asr_name = args["--asr"]
+    whisper_name = args["--whisper"]
+    debug = args["--debug"]
+    show = args["--show"]
 
     from os.path import exists
 
@@ -89,26 +90,26 @@ def main() -> None:
     from sonusai.utils import write_audio
 
     ts = create_timestamp()
-    capture_name = f'audiofe_capture_{ts}'
-    capture_wav = capture_name + '.wav'
-    capture_png = capture_name + '.png'
-    predict_name = f'audiofe_predict_{ts}'
-    predict_wav = predict_name + '.wav'
-    predict_png = predict_name + '.png'
-    h5_name = f'audiofe_{ts}.h5'
+    capture_name = f"audiofe_capture_{ts}"
+    capture_wav = capture_name + ".wav"
+    capture_png = capture_name + ".png"
+    predict_name = f"audiofe_predict_{ts}"
+    predict_wav = predict_name + ".wav"
+    predict_png = predict_name + ".png"
+    h5_name = f"audiofe_{ts}.h5"
 
     # Setup logging file
-    create_file_handler('audiofe.log')
+    create_file_handler("audiofe.log")
     update_console_handler(verbose)
-    initial_log_messages('audiofe')
+    initial_log_messages("audiofe")
 
     if show:
-        logger.info('List of available audio inputs:')
-        logger.info('')
+        logger.info("List of available audio inputs:")
+        logger.info("")
         p = pyaudio.PyAudio()
         for name in get_input_devices(p):
-            logger.info(f'{name}')
-        logger.info('')
+            logger.info(f"{name}")
+        logger.info("")
         p.terminate()
         return
 
@@ -122,27 +123,27 @@ def main() -> None:
             return
         # Only write if capture from device, not for file input
         write_audio(capture_wav, capture_audio, SAMPLE_RATE)
-        logger.info('')
-        logger.info(f'Wrote capture audio with shape {capture_audio.shape} to {capture_wav}')
+        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 feature data with shape {capture_audio.shape} to {h5_name}')
+        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 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} ...')
+        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}')
+        logger.info(f"Capture audio ASR: {capture_asr}")
 
     if model_name is not None:
         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.')
+            logger.error("Error: ONNX model does not have required SonusAI hyperparameters, cannot proceed.")
             raise SystemExit(1)
-        feature_mode = hparams['feature']
+        feature_mode = hparams["feature"]
         in0name = sess_inputs[0].name
         in0type = sess_inputs[0].type
         out_names = [n.name for n in session.get_outputs()]
@@ -150,47 +151,50 @@ def main() -> None:
         # 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}')
+        logger.info(f"Wrote capture plots to {capture_png}")
 
         if debug:
-            with h5py.File(h5_name, 'a') as f:
-                if 'feature' in f:
-                    del f['feature']
-                f.create_dataset('feature', data=feature)
-            logger.info(f'Wrote feature with shape {feature.shape} to {h5_name}')
+            with h5py.File(h5_name, "a") as f:
+                if "feature" in f:
+                    del f["feature"]
+                f.create_dataset("feature", data=feature)
+            logger.info(f"Wrote feature with shape {feature.shape} to {h5_name}")
 
-        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
+        if in0type.find("float16") != -1:
+            logger.info("Detected input of float16, converting all feature inputs to that type.")
+            feature = np.float16(feature)  # type: ignore[assignment]
 
         # 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))
+        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:
-                    del f['predict']
-                f.create_dataset('predict', data=predict)
-            logger.info(f'Wrote predict with shape {predict.shape} to {h5_name}')
+            with h5py.File(h5_name, "a") as f:
+                if "predict" in f:
+                    del f["predict"]
+                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=feature_mode)
         write_audio(predict_wav, predict_audio, SAMPLE_RATE)
-        logger.info(f'Wrote predict audio with shape {predict_audio.shape} to {predict_wav}')
+        logger.info(f"Wrote predict audio with shape {predict_audio.shape} to {predict_wav}")
         if debug:
-            with h5py.File(h5_name, 'a') as f:
-                if 'predict_audio' in f:
-                    del f['predict_audio']
-                f.create_dataset('predict_audio', data=predict_audio)
-            logger.info(f'Wrote predict audio with shape {predict_audio.shape} to {h5_name}')
+            with h5py.File(h5_name, "a") as f:
+                if "predict_audio" in f:
+                    del f["predict_audio"]
+                f.create_dataset("predict_audio", data=predict_audio)
+            logger.info(f"Wrote predict audio with shape {predict_audio.shape} to {h5_name}")
 
         save_figure(predict_png, predict_audio, predict)
-        logger.info(f'Wrote predict plots to {predict_png}')
+        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} ...')
+            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}')
+            logger.info(f"Predict audio ASR: {predict_asr}")
 
 
 def get_frames_from_device(input_name: str | None, length: float, chunk: int = 1024) -> AudioT:
@@ -209,32 +213,34 @@ def get_frames_from_device(input_name: str | None, length: float, chunk: int = 1
 
     input_devices = get_input_devices(p)
     if not input_devices:
-        raise ValueError('No input audio devices found')
+        raise ValueError("No input audio devices found")
 
     if input_name is None:
         input_name = input_devices[0]
 
     try:
         device_index = get_input_device_index_by_name(p, input_name)
-    except ValueError:
-        msg = f'Could not find {input_name}\n'
-        msg += f'Available devices:\n'
+    except ValueError as e:
+        msg = f"Could not find {input_name}\n"
+        msg += "Available devices:\n"
         for input_device in input_devices:
-            msg += f'  {input_device}\n'
-        raise ValueError(msg)
-
-    logger.info(f'Capturing from {p.get_device_info_by_index(device_index).get("name")}')
-    stream = p.open(format=pyaudio.paFloat32,
-                    channels=CHANNEL_COUNT,
-                    rate=SAMPLE_RATE,
-                    input=True,
-                    input_device_index=device_index)
+            msg += f"  {input_device}\n"
+        raise ValueError(msg) from e
+
+    logger.info(f"Capturing from {p.get_device_info_by_index(device_index).get('name')}")
+    stream = p.open(
+        format=pyaudio.paFloat32,
+        channels=CHANNEL_COUNT,
+        rate=SAMPLE_RATE,
+        input=True,
+        input_device_index=device_index,
+    )
     stream.start_stream()
 
     print()
-    print('+---------------------------------+')
-    print('| Press Enter to stop             |')
-    print('+---------------------------------+')
+    print("+---------------------------------+")
+    print("| Press Enter to stop             |")
+    print("+---------------------------------+")
     print()
 
     elapsed = 0.0
@@ -243,14 +249,21 @@ def get_frames_from_device(input_name: str | None, length: float, chunk: int = 1
     while elapsed < length or length == -1:
         raw_frames.append(stream.read(num_frames=chunk, exception_on_overflow=False))
         elapsed += seconds_per_chunk
-        if select([stdin, ], [], [], 0)[0]:
+        if select(
+            [
+                stdin,
+            ],
+            [],
+            [],
+            0,
+        )[0]:
             stdin.read(1)
             length = elapsed
 
     stream.stop_stream()
     stream.close()
     p.terminate()
-    frames = np.frombuffer(b''.join(raw_frames), dtype=np.float32)
+    frames = np.frombuffer(b"".join(raw_frames), dtype=np.float32)
     return frames
 
 
@@ -259,7 +272,7 @@ def get_frames_from_file(input_name: str, length: float) -> AudioT:
     from sonusai.mixture import SAMPLE_RATE
     from sonusai.mixture import read_audio
 
-    logger.info(f'Capturing from {input_name}')
+    logger.info(f"Capturing from {input_name}")
     frames = read_audio(input_name)
     if length != -1:
         num_frames = int(length * SAMPLE_RATE)
@@ -289,16 +302,16 @@ def save_figure(name: str, audio: np.ndarray, feature: np.ndarray) -> None:
     fig, (ax1, ax2) = plt.subplots(nrows=2)
     ax1.set_title(name)
     ax1.plot(t, audio[:samples])
-    ax1.set_ylabel('Signal')
+    ax1.set_ylabel("Signal")
     ax1.set_xlim(0, length_in_s)
     ax1.set_ylim(-1, 1)
 
-    ax2.imshow(spectrum, origin='lower', aspect='auto')
+    ax2.imshow(spectrum, origin="lower", aspect="auto")
     ax2.set_xticks([])
-    ax2.set_ylabel('Feature')
+    ax2.set_ylabel("Feature")
 
     plt.savefig(name, dpi=300)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()