Update preprocessing

README.md

@@ -105,6 +105,15 @@ class PowerToDB(torch.nn.Module):
 
         return log_spec
 
+# Initialize preprocessors
+spectrogram_converter = torchaudio.transforms.Spectrogram(
+    n_fft=2048, hop_length=256, power=2.0
+)
+mel_converter = torchaudio.transforms.MelScale(
+    n_mels=256, n_stft=1025, sample_rate=32_000
+)
+powerToDB = PowerToDB(top_db=80)
+
 
 def preprocess(audio, sample_rate_of_audio):
     """
@@ -114,32 +123,26 @@ def preprocess(audio, sample_rate_of_audio):
     - Normalize the melscale spectrogram with mean: -4.268, std: 4.569 (from AudioSet)
 
     """
-    powerToDB = PowerToDB()
-    # Resample to 32kHz
-    resample = torchaudio.transforms.Resample(
-        orig_freq=sample_rate_of_audio, new_freq=32000
-    )
-    audio = resample(audio)
-    spectrogram = torchaudio.transforms.Spectrogram(
-        n_fft=2048, hop_length=256, power=2.0
-    )(audio)
-    melspec = torchaudio.transforms.MelScale(n_mels=256, n_stft=1025)(spectrogram)
+    spectrogram = spectrogram_converter(audio)
+    spectrogram = spectrogram.to(torch.float32)
+    melspec = mel_converter(spectrogram)
     dbscale = powerToDB(melspec)
     normalized_dbscale = transforms.Normalize((-4.268,), (4.569,))(dbscale)
+    # add batch dimension if needed
+    if normalized_dbscale.dim() == 3:
+        normalized_dbscale = normalized_dbscale.unsqueeze(0)
     return normalized_dbscale
 
 preprocessed_audio = preprocess(audio, sample_rate)
 print("Preprocessed_audio shape:", preprocessed_audio.shape)
 
-
-logits = model(preprocessed_audio.unsqueeze(0)).logits
+logits = model(preprocessed_audio).logits
 print("Logits shape: ", logits.shape)
 
 top5 = torch.topk(logits, 5)
 print("Top 5 logits:", top5.values)
 print("Top 5 predicted classes:")
 print([model.config.id2label[i] for i in top5.indices.squeeze().tolist()])
-
 ```
 
 ## Model Source