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

sonusai/{plot.py → deprecated/plot.py}

@@ -57,19 +57,21 @@ 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)
 
 
 signal.signal(signal.SIGINT, signal_handler)
 
 
-def spec_plot(mixture: AudioT,
-              feature: Feature,
-              predict: Predict = None,
-              target: AudioT = None,
-              labels: list[str] = None,
-              title: str = '') -> plt.Figure:
+def spec_plot(
+    mixture: AudioT,
+    feature: Feature,
+    predict: Predict | None = None,
+    target: AudioT | None = None,
+    labels: list[str] | None = None,
+    title: str = "",
+) -> plt.Figure:
     from sonusai.mixture import SAMPLE_RATE
 
     num_plots = 4 if predict is not None else 2
@@ -77,46 +79,53 @@ def spec_plot(mixture: AudioT,
 
     # Plot the waveform
     x_axis = np.arange(len(mixture), dtype=np.float32) / SAMPLE_RATE
-    ax[0].plot(x_axis, mixture, label='Mixture')
+    ax[0].plot(x_axis, mixture, label="Mixture")
     ax[0].set_xlim(x_axis[0], x_axis[-1])
     ax[0].set_ylim([-1.025, 1.025])
     if target is not None:
         # Plot target time-domain waveform on top of mixture
-        color = 'tab:blue'
-        ax[0].plot(x_axis, target, color=color, label='Target')
-        ax[0].set_ylabel('magnitude', color=color)
-    ax[0].set_title('Waveform')
+        color = "tab:blue"
+        ax[0].plot(x_axis, target, color=color, label="Target")
+        ax[0].set_ylabel("magnitude", color=color)
+    ax[0].set_title("Waveform")
 
     # Plot the spectrogram
-    ax[1].imshow(np.transpose(feature), aspect='auto', interpolation='nearest', origin='lower')
-    ax[1].set_title('Feature')
+    ax[1].imshow(np.transpose(feature), aspect="auto", interpolation="nearest", origin="lower")
+    ax[1].set_title("Feature")
 
     if predict is not None:
+        if labels is None:
+            raise ValueError("Provided predict without labels")
+
         # Plot and label the model output scores for the top-scoring classes.
         mean_predict = np.mean(predict, axis=0)
         num_classes = predict.shape[-1]
         top_n = min(10, num_classes)
         top_class_indices = np.argsort(mean_predict)[::-1][:top_n]
-        ax[2].imshow(np.transpose(predict[:, top_class_indices]), aspect='auto', interpolation='nearest', cmap='gray_r')
+        ax[2].imshow(
+            np.transpose(predict[:, top_class_indices]),
+            aspect="auto",
+            interpolation="nearest",
+            cmap="gray_r",
+        )
         y_ticks = range(0, top_n)
         ax[2].set_yticks(y_ticks, [labels[top_class_indices[x]] for x in y_ticks])
         ax[2].set_ylim(-0.5 + np.array([top_n, 0]))
-        ax[2].set_title('Class Scores')
+        ax[2].set_title("Class Scores")
 
         # Plot the probabilities
         ax[3].plot(predict[:, top_class_indices])
-        ax[3].legend(np.array(labels)[top_class_indices], loc='best')
-        ax[3].set_title('Class Probabilities')
+        ax[3].legend(np.array(labels)[top_class_indices], loc="best")
+        ax[3].set_title("Class Probabilities")
 
     fig.suptitle(title)
 
     return fig
 
 
-def spec_energy_plot(mixture: AudioT,
-                     feature: Feature,
-                     truth_f: Truth = None,
-                     predict: Predict = None) -> plt.Figure:
+def spec_energy_plot(
+    mixture: AudioT, feature: Feature, truth_f: Truth | None = None, predict: Predict | None = None
+) -> plt.Figure:
     from sonusai.mixture import SAMPLE_RATE
 
     num_plots = 2
@@ -130,35 +139,47 @@ def spec_energy_plot(mixture: AudioT,
     # Plot the waveform
     p = 0
     x_axis = np.arange(len(mixture), dtype=np.float32) / SAMPLE_RATE
-    ax[p].plot(x_axis, mixture, label='Mixture')
+    ax[p].plot(x_axis, mixture, label="Mixture")
     ax[p].set_xlim(x_axis[0], x_axis[-1])
     ax[p].set_ylim([-1.025, 1.025])
-    ax[p].set_title('Waveform')
+    ax[p].set_title("Waveform")
 
     # Plot the spectrogram
     p += 1
-    ax[p].imshow(np.transpose(feature), aspect='auto', interpolation='nearest', origin='lower')
-    ax[p].set_title('Feature')
+    ax[p].imshow(np.transpose(feature), aspect="auto", interpolation="nearest", origin="lower")
+    ax[p].set_title("Feature")
 
     if truth_f is not None:
         p += 1
-        ax[p].imshow(np.transpose(truth_f), aspect='auto', interpolation='nearest', origin='lower')
-        ax[p].set_title('Truth')
+        ax[p].imshow(
+            np.transpose(truth_f),
+            aspect="auto",
+            interpolation="nearest",
+            origin="lower",
+        )
+        ax[p].set_title("Truth")
 
     if predict is not None:
         p += 1
-        ax[p].imshow(np.transpose(predict), aspect='auto', interpolation='nearest', origin='lower')
-        ax[p].set_title('Predict')
+        ax[p].imshow(
+            np.transpose(predict),
+            aspect="auto",
+            interpolation="nearest",
+            origin="lower",
+        )
+        ax[p].set_title("Predict")
 
     return fig
 
 
-def class_plot(mixture: AudioT,
-               target: AudioT = None,
-               truth_f: Truth = None,
-               predict: Predict = None,
-               label: str = '') -> plt.Figure:
-    """ Plot mixture waveform with optional prediction and/or truth together in a single plot
+def class_plot(
+    mixture: AudioT,
+    target: AudioT | None = None,
+    truth_f: Truth | None = None,
+    predict: Predict | None = None,
+    label: str = "",
+) -> plt.Figure:
+    """Plot mixture waveform with optional prediction and/or truth together in a single plot
 
     The target waveform can optionally be provided, and prediction and truth can have multiple classes.
 
@@ -174,30 +195,30 @@ def class_plot(mixture: AudioT,
     from sonusai.mixture import SAMPLE_RATE
 
     if mixture.ndim != 1:
-        raise SonusAIError('Too many dimensions in mixture')
+        raise SonusAIError("Too many dimensions in mixture")
 
     if target is not None and target.ndim != 1:
-        raise SonusAIError('Too many dimensions in target')
+        raise SonusAIError("Too many dimensions in target")
 
     # Set default to 1 frame when there is no truth or predict data
     frames = 1
     if truth_f is not None and predict is not None:
         if truth_f.ndim != 1:
-            raise SonusAIError('Too many dimensions in truth_f')
+            raise SonusAIError("Too many dimensions in truth_f")
         t_frames = len(truth_f)
 
         if predict.ndim != 1:
-            raise SonusAIError('Too many dimensions in predict')
+            raise SonusAIError("Too many dimensions in predict")
         p_frames = len(predict)
 
         frames = min(t_frames, p_frames)
     elif truth_f is not None:
         if truth_f.ndim != 1:
-            raise SonusAIError('Too many dimensions in truth_f')
+            raise SonusAIError("Too many dimensions in truth_f")
         frames = len(truth_f)
     elif predict is not None:
         if predict.ndim != 1:
-            raise SonusAIError('Too many dimensions in predict')
+            raise SonusAIError("Too many dimensions in predict")
         frames = len(predict)
 
     samples = (len(mixture) // frames) * frames
@@ -208,41 +229,51 @@ def class_plot(mixture: AudioT,
     fig, ax = plt.subplots(1, 1, constrained_layout=True, figsize=(11, 8.5))
 
     # Plot the time-domain waveforms then truth/prediction on second axis
-    ax.plot(x_axis, mixture[0:samples], color='mistyrose', label='Mixture')
-    color = 'red'
+    ax.plot(x_axis, mixture[0:samples], color="mistyrose", label="Mixture")
+    color = "red"
     ax.set_xlim(x_axis[0], x_axis[-1])
-    ax.set_ylim([-1.025, 1.025])
-    ax.set_ylabel(f'Amplitude', color=color)
-    ax.tick_params(axis='y', labelcolor=color)
+    ax.set_ylim((-1.025, 1.025))
+    ax.set_ylabel("Amplitude", color=color)
+    ax.tick_params(axis="y", labelcolor=color)
 
     # Plot target time-domain waveform
     if target is not None:
-        ax.plot(x_axis, target[0:samples], color='blue', label='Target')
+        ax.plot(x_axis, target[0:samples], color="blue", label="Target")
 
     # instantiate 2nd y-axis that shares the same x-axis
     if truth_f is not None or predict is not None:
-        y_label = 'Truth/Predict'
+        y_label = "Truth/Predict"
         if truth_f is None:
-            y_label = 'Predict'
+            y_label = "Predict"
         if predict is None:
-            y_label = 'Truth'
+            y_label = "Truth"
 
         ax2 = ax.twinx()
 
-        color = 'black'
+        color = "black"
         ax2.set_xlim(x_axis[0], x_axis[-1])
-        ax2.set_ylim([-0.025, 1.025])
+        ax2.set_ylim((-0.025, 1.025))
         ax2.set_ylabel(y_label, color=color)
-        ax2.tick_params(axis='y', labelcolor=color)
+        ax2.tick_params(axis="y", labelcolor=color)
 
         if truth_f is not None:
-            ax2.plot(x_axis, expand_frames_to_samples(truth_f, samples), color='green', label='Truth')
+            ax2.plot(
+                x_axis,
+                expand_frames_to_samples(truth_f, samples),
+                color="green",
+                label="Truth",
+            )
 
         if predict is not None:
-            ax2.plot(x_axis, expand_frames_to_samples(predict, samples), color='brown', label='Predict')
+            ax2.plot(
+                x_axis,
+                expand_frames_to_samples(predict, samples),
+                color="brown",
+                label="Predict",
+            )
 
     # set only on last/bottom plot
-    ax.set_xlabel('time (s)')
+    ax.set_xlabel("time (s)")
 
     fig.suptitle(label)
 
@@ -263,7 +294,6 @@ def main() -> None:
     args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
 
     from dataclasses import asdict
-
     from os.path import basename
     from os.path import exists
     from os.path import isdir
@@ -279,117 +309,121 @@ def main() -> None:
     from sonusai import initial_log_messages
     from sonusai import logger
     from sonusai import update_console_handler
-    from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import FeatureGeneratorConfig
+    from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import get_feature_from_audio
     from sonusai.mixture import get_truth_indices_for_mixid
     from sonusai.mixture import read_audio
     from sonusai.utils import get_label_names
     from sonusai.utils import print_mixture_details
 
-    verbose = args['--verbose']
-    model_name = args['--model']
-    output_name = args['--output']
-    labels_name = args['--labels']
-    mixid = args['--mixid']
-    energy = args['--energy']
-    input_name = args['INPUT']
+    verbose = args["--verbose"]
+    model_name = args["--model"]
+    output_name = args["--output"]
+    labels_name = args["--labels"]
+    mixid = args["--mixid"]
+    energy = args["--energy"]
+    input_name = args["INPUT"]
 
     if mixid is not None:
         mixid = int(mixid)
 
-    create_file_handler('plot.log')
+    create_file_handler("plot.log")
     update_console_handler(verbose)
-    initial_log_messages('plot')
+    initial_log_messages("plot")
 
     if not exists(input_name):
-        raise SonusAIError(f'{input_name} does not exist')
+        raise SonusAIError(f"{input_name} does not exist")
 
-    logger.info('')
-    logger.info(f'Input:  {input_name}')
+    logger.info("")
+    logger.info(f"Input:  {input_name}")
     if model_name is not None:
-        logger.info(f'Model:  {model_name}')
+        logger.info(f"Model:  {model_name}")
     if output_name is not None:
-        logger.info(f'Output: {output_name}')
-    logger.info('')
+        logger.info(f"Output: {output_name}")
+    logger.info("")
 
     ext = splitext(input_name)[1]
 
     model = None
     target_audio = None
     truth_f = None
-    t_indices = None
+    t_indices = []
 
     if model_name is not None:
         model = Predict(model_name)
 
-    if ext == '.wav':
-        if model_name is None:
-            raise SonusAIError('Must specify MODEL when input is WAV')
+    if ext == ".wav":
+        if model is None:
+            raise SonusAIError("Must specify MODEL when input is WAV")
 
         mixture_audio = read_audio(input_name)
         feature = get_feature_from_audio(audio=mixture_audio, feature_mode=model.feature)
-        fg_config = FeatureGeneratorConfig(feature_mode=model.feature,
-                                           num_classes=model.output_shape[-1],
-                                           truth_mutex=False)
+        fg_config = FeatureGeneratorConfig(
+            feature_mode=model.feature,
+            num_classes=model.output_shape[-1],
+            truth_mutex=False,
+        )
         fg = FeatureGenerator(**asdict(fg_config))
         fg_step = fg.step
         mixdb = None
-        logger.debug(f'Audio samples      {len(mixture_audio)}')
-        logger.debug(f'Feature shape      {feature.shape}')
+        logger.debug(f"Audio samples      {len(mixture_audio)}")
+        logger.debug(f"Feature shape      {feature.shape}")
 
     elif isdir(input_name):
         if mixid is None:
-            raise SonusAIError('Must specify mixid when input is mixture database')
+            raise SonusAIError("Must specify mixid when input is mixture database")
 
         mixdb = MixtureDatabase(input_name)
         fg_step = mixdb.fg_step
 
         print_mixture_details(mixdb=mixdb, mixid=mixid, desc_len=24, print_fn=logger.info)
 
-        logger.info(f'Generating data for mixture {mixid}')
+        logger.info(f"Generating data for mixture {mixid}")
         mixture_audio = mixdb.mixture_mixture(mixid)
         target_audio = mixdb.mixture_target(mixid)
         feature, truth_f = mixdb.mixture_ft(mixid)
         t_indices = [x - 1 for x in get_truth_indices_for_mixid(mixdb=mixdb, mixid=mixid)]
 
         target_files = [mixdb.target_file(target.file_id) for target in mixdb.mixtures[mixid].targets]
-        truth_functions = list(set([sub2.function for sub1 in target_files for sub2 in sub1.truth_settings]))
-        energy = 'energy_f' in truth_functions or 'snr_f' in truth_functions
+        truth_functions = list({sub2.function for sub1 in target_files for sub2 in sub1.truth_configs})
+        energy = "energy_f" in truth_functions or "snr_f" in truth_functions
 
-        logger.debug(f'Audio samples      {len(mixture_audio)}')
-        logger.debug(f'Targets:')
+        logger.debug(f"Audio samples      {len(mixture_audio)}")
+        logger.debug("Targets:")
         mixture = mixdb.mixture(mixid)
         for target in mixture.targets:
             target_file = mixdb.target_file(target.file_id)
             name = target_file.name
             duration = target_file.duration
             augmentation = target.augmentation
-            logger.debug(f'  Name             {name}')
-            logger.debug(f'  Duration         {duration}')
-            logger.debug(f'  Augmentation     {augmentation}')
+            logger.debug(f"  Name             {name}")
+            logger.debug(f"  Duration         {duration}")
+            logger.debug(f"  Augmentation     {augmentation}")
 
-        logger.debug(f'Feature shape      {feature.shape}')
-        logger.debug(f'Truth shape        {truth_f.shape}')
+        logger.debug(f"Feature shape      {feature.shape}")
+        logger.debug(f"Truth shape        {truth_f.shape}")
 
     else:
-        raise SonusAIError(f'Unknown file type for {input_name}')
+        raise SonusAIError(f"Unknown file type for {input_name}")
 
     predict = None
     labels = None
     indices = []
     if model is not None:
-        logger.debug('')
-        logger.info(f'Running prediction on mixture {mixid}')
-        logger.debug(f'Model feature name {model.feature}')
-        logger.debug(f'Model input shape  {model.input_shape}')
-        logger.debug(f'Model output shape {model.output_shape}')
+        logger.debug("")
+        logger.info(f"Running prediction on mixture {mixid}")
+        logger.debug(f"Model feature name {model.feature}")
+        logger.debug(f"Model input shape  {model.input_shape}")
+        logger.debug(f"Model output shape {model.output_shape}")
 
         if feature.shape[0] < model.input_shape[0]:
-            raise SonusAIError(f'Mixture {mixid} contains {feature.shape[0]} '
-                               f'frames of data which is not enough to run prediction; '
-                               f'at least {model.input_shape[0]} frames are needed for this model.\n'
-                               f'Consider using a model with a smaller batch size or a mixture with more data.')
+            raise SonusAIError(
+                f"Mixture {mixid} contains {feature.shape[0]} "
+                f"frames of data which is not enough to run prediction; "
+                f"at least {model.input_shape[0]} frames are needed for this model.\n"
+                f"Consider using a model with a smaller batch size or a mixture with more data."
+            )
 
         predict = model.execute(feature)
 
@@ -400,10 +434,10 @@ def main() -> None:
         p_indices = np.argsort(p_max)[::-1][:5]
         p_max_len = max([len(labels[i]) for i in p_indices])
 
-        logger.info('Top 5 active prediction classes by max:')
+        logger.info("Top 5 active prediction classes by max:")
         for p_index in p_indices:
-            logger.info(f'  {labels[p_index]:{p_max_len}s} {p_max[p_index]:.3f}')
-        logger.info('')
+            logger.info(f"  {labels[p_index]:{p_max_len}s} {p_max[p_index]:.3f}")
+        logger.info("")
 
         indices = list(p_indices)
 
@@ -414,26 +448,30 @@ def main() -> None:
 
     base_name = basename(splitext(input_name)[0])
     if mixdb is not None:
-        title = f'{input_name} Mixture {mixid}'
-        pdf_name = f'{base_name}-mix{mixid}-plot.pdf'
+        title = f"{input_name} Mixture {mixid}"
+        pdf_name = f"{base_name}-mix{mixid}-plot.pdf"
     else:
-        title = f'{input_name}'
-        pdf_name = f'{base_name}-plot.pdf'
+        title = f"{input_name}"
+        pdf_name = f"{base_name}-plot.pdf"
 
     # Original size [frames, stride, feature_parameters]
     # Decimate in the stride dimension
     # Reshape to get frames*decimated_stride, feature_parameters
     if feature.ndim != 3:
-        raise SonusAIError(f'feature does not have 3 dimensions: frames, stride, feature_parameters')
+        raise SonusAIError("feature does not have 3 dimensions: frames, stride, feature_parameters")
     spectrogram = feature[:, -fg_step:, :]
     spectrogram = np.reshape(spectrogram, (spectrogram.shape[0] * spectrogram.shape[1], spectrogram.shape[2]))
 
     with PdfPages(pdf_name) as pdf:
-        pdf.savefig(spec_plot(mixture=mixture_audio,
-                              feature=spectrogram,
-                              predict=predict,
-                              labels=labels,
-                              title=title))
+        pdf.savefig(
+            spec_plot(
+                mixture=mixture_audio,
+                feature=spectrogram,
+                predict=predict,
+                labels=labels,
+                title=title,
+            )
+        )
         for index in indices:
             if energy:
                 t_tmp = None
@@ -444,10 +482,14 @@ def main() -> None:
                 if predict is not None:
                     p_tmp = 10 * np.log10(predict + np.finfo(np.float32).eps)
 
-                pdf.savefig(spec_energy_plot(mixture=mixture_audio,
-                                             feature=spectrogram,
-                                             truth_f=t_tmp,
-                                             predict=p_tmp))
+                pdf.savefig(
+                    spec_energy_plot(
+                        mixture=mixture_audio,
+                        feature=spectrogram,
+                        truth_f=t_tmp,
+                        predict=p_tmp,
+                    )
+                )
             else:
                 p_tmp = None
                 if predict is not None:
@@ -457,18 +499,22 @@ def main() -> None:
                 if labels is not None:
                     l_tmp = labels[index]
 
-                pdf.savefig(class_plot(mixture=mixture_audio,
-                                       target=target_audio[index],
-                                       truth_f=truth_f[:, index],
-                                       predict=p_tmp,
-                                       label=l_tmp))
-        logger.info(f'Wrote {pdf_name}')
+                pdf.savefig(
+                    class_plot(
+                        mixture=mixture_audio,
+                        target=target_audio[index],
+                        truth_f=truth_f[:, index],
+                        predict=p_tmp,
+                        label=l_tmp,
+                    )
+                )
+        logger.info(f"Wrote {pdf_name}")
 
     if output_name:
-        with h5py.File(output_name, 'w') as f:
-            f.create_dataset(name='predict', data=predict)
-            logger.info(f'Wrote {output_name}')
+        with h5py.File(output_name, "w") as f:
+            f.create_dataset(name="predict", data=predict)
+            logger.info(f"Wrote {output_name}")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()