small formatting fixes

Author: saltytine

basic_pitch/callbacks.py

@@ -17,14 +17,15 @@
 
 import os
 
-from typing import Any, Dict
+from typing import Any, Dict, Callable
 
 import tensorflow as tf
 
 from basic_pitch import visualize
 
 TENSORBOARD_LOGS_SUBDIR = "tensorboard_logs"
 
+
 class VisualizeCallback(tf.keras.callbacks.Callback):
     """
     Callback to run during training to create tensorboard visualizations per epoch.
@@ -49,10 +50,10 @@ def __init__(
         use_tf_function: bool = True,
     ):
         super().__init__()
-        self.train_ds = train_ds.take(max_batches).prefatch(prefetch_batches)
+        self.train_ds = train_ds.take(max_batches).prefetch(prefetch_batches)
         self.validation_ds = validation_ds.take(max_batches).prefetch(prefetch_batches)
         self.tensorboard_dir = os.path.join(tensorboard_dir, TENSORBOARD_LOGS_SUBDIR)
-        self.file_writer = tf.summary.create_file_writer(tensorboard_dir)
+        self.file_writer = tf.summary.create_file_writer(self.tensorboard_dir)
         self.sonify = sonify
         self.contours = contours
         self.use_tf_function = use_tf_function
@@ -62,20 +63,18 @@ def __init__(
 
         self._predict_fn = None
 
-    def set_module(self, model):
+    def set_model(self, model: tf.keras.Model) -> None:
         super().set_model(model)
         if self.use_tf_function:
-            @tf_function
-            def fast_predict(inputs):
+            def fast_predict(inputs: tf.Tensor) -> Any:
                 return model(inputs, training=False)
-            self._predict_fn = fast_predict
+            self._predict_fn = tf.function(fast_predict)
         else:
             self._predict_fn = model.predict
 
-    def _predict(self, inputs):
+    def _predict(self, inputs: tf.Tensor) -> Any:
         if self._predict_fn is not None:
             outputs = self._predict_fn(inputs)
-            # tf.functions might output as a dict of tensors, convert to numpy:
             if isinstance(outputs, dict):
                 outputs = {k: v.numpy() if hasattr(v, "numpy") else v for k, v in outputs.items()}
             return outputs
@@ -90,15 +89,16 @@ def on_epoch_end(self, epoch: int, logs: Dict[Any, Any]) -> None:
             for batch in ds:
                 inputs, targets = batch[:2]
                 outputs = self._predict(inputs)
+                loss_val = logs.get(loss_key)
                 visualize.visualize_transcription(
                     self.file_writer,
                     stage,
                     inputs,
                     targets,
                     outputs,
-                    logs.get(loss_key),
+                    float(loss_val) if loss_val is not None else 0.0,
                     epoch,
                     sonify=self.sonify,
-                    contours=self.contours
+                    contours=self.contours,
                 )
-                break
+                break

basic_pitch/commandline_printing.py

@@ -23,6 +23,8 @@
 
 TF_LOG_LEVEL_KEY = "TF_CPP_MIN_LOG_LEVEL"
 TF_LOG_LEVEL_NO_WARNINGS_VALUE = "3"
+DEFAULT_PRINT_INDENT = "  "
+
 s_print_lock = threading.Lock()
 OUTPUT_EMOJIS = {
     "MIDI": "💅",

basic_pitch/constants.py

@@ -19,6 +19,7 @@
 
 from enum import Enum
 
+
 SEMITONES_PER_OCTAVE = 12 # for frequency bin calculations
 
 FFT_HOP = 256

basic_pitch/layers/nnaudio.py

@@ -41,6 +41,7 @@
 DEFAULT_CQT_OUTPUT_FORMAT = "Magnitude"
 DEFAULT_LOW_PASS_TRANSITION_BANDWIDTH = 0.001
 
+
 def create_lowpass_filter(
     band_center: float = DEFAULT_BAND_CENTER,
     kernel_length: int = DEFAULT_KERNEL_LENGTH,

basic_pitch/models.py

@@ -51,6 +51,7 @@
 ONSET_STRIDES_1 = (1, 3)
 ONSET_KERNEL_SIZE_2 = (3, 3)
 
+
 def transcription_loss(y_true: tf.Tensor, y_pred: tf.Tensor, label_smoothing: float) -> tf.Tensor:
     """Really a binary cross entropy loss. Used to calculate the loss between the predicted
     posteriorgrams and the ground truth matrices.

basic_pitch/nn.py

@@ -24,6 +24,7 @@
 
 SEMITONES_PER_OCTAVE = 12
 
+
 class HarmonicStacking(tf.keras.layers.Layer):
     """Harmonic stacking layer
 

basic_pitch/note_creation.py

@@ -48,6 +48,7 @@
 MIDI_VELOCITY_SCALE = 127
 PITCH_BEND_SCALE = 4096
 
+
 def model_output_to_notes(
     output: Dict[str, np.array],
     onset_thresh: float,