From f36bab95ddf250043e3e2ec776eaf9959ec19e4f Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Wed, 10 Nov 2021 20:06:33 -0700
Subject: [PATCH] Audio resample injector

---
 codes/trainer/injectors/base_injectors.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/codes/trainer/injectors/base_injectors.py b/codes/trainer/injectors/base_injectors.py
index 1e11bf72..4c05384d 100644
--- a/codes/trainer/injectors/base_injectors.py
+++ b/codes/trainer/injectors/base_injectors.py
@@ -1,6 +1,7 @@
 import random
 
 import torch.nn
+import torchaudio.functional
 from kornia.augmentation import RandomResizedCrop
 from torch.cuda.amp import autocast
 
@@ -569,6 +570,17 @@ class RandomAudioCropInjector(Injector):
         return {self.output: inp[:, :, start:start+self.crop_sz]}
 
 
+class AudioResampleInjector(Injector):
+    def __init__(self, opt, env):
+        super().__init__(opt, env)
+        self.input_sr = opt['input_sample_rate']
+        self.output_sr = opt['output_sample_rate']
+
+    def forward(self, state):
+        inp = state[self.input]
+        return {self.output: torchaudio.functional.resample(inp, self.input_sr, self.output_sr)}
+
+
 if __name__ == '__main__':
     inj = MelSpectrogramInjector({'in': 'x', 'out': 'y'}, None)
     print(inj({'x':torch.rand(10,1,40800)})['y'].shape)
\ No newline at end of file