Update pipelines/data/data_module.py

pipelines/data/data_module.py

@@ -5,6 +5,7 @@
 # Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
 
 import torch
+import whisper
 import torchaudio
 import torchvision
 from .transforms import AudioTransform, VideoTransform
@@ -28,8 +29,9 @@ class AVSRDataLoader:
 
     def load_data(self, data_filename, landmarks=None, transform=True):
         if self.modality == "audio":
-            audio, sample_rate = self.load_audio(data_filename)
-            audio = self.audio_process(audio, sample_rate)
+            # audio, sample_rate = self.load_audio(data_filename)
+            # audio = self.audio_process(audio, sample_rate)
+            audio = self.load_audio(data_filename)
             return self.audio_transform(audio) if self.transform else audio
         if self.modality == "video":
             video = self.load_video(data_filename)
@@ -38,8 +40,9 @@ class AVSRDataLoader:
             return self.video_transform(video) if self.transform else video
         if self.modality == "audiovisual":
             rate_ratio = 640
-            audio, sample_rate = self.load_audio(data_filename)
-            audio = self.audio_process(audio, sample_rate)
+            # audio, sample_rate = self.load_audio(data_filename)
+            # audio = self.audio_process(audio, sample_rate)
+            audio = self.load_audio(data_filename)
             video = self.load_video(data_filename)
             video = self.video_process(video, landmarks)
             video = torch.tensor(video)
@@ -53,16 +56,18 @@ class AVSRDataLoader:
 
 
     def load_audio(self, data_filename):
-        waveform, sample_rate = torchaudio.load(data_filename, normalize=True)
-        return waveform, sample_rate
+        # rtype: [1, T]
+        waveform = torch.tensor(whisper.load_audio(data_filename)).unsqueeze(0)
+        # waveform, sample_rate = torchaudio.load(data_filename, normalize=True)
+        # return waveform, sample_rate
 
 
     def load_video(self, data_filename):
         return torchvision.io.read_video(data_filename, pts_unit='sec')[0].numpy()
 
 
-    def audio_process(self, waveform, sample_rate, target_sample_rate=16000):
-        if sample_rate != target_sample_rate:
-            waveform = torchaudio.functional.resample(waveform, sample_rate, target_sample_rate)
-        waveform = torch.mean(waveform, dim=0, keepdim=True)
-        return waveform
+    # def audio_process(self, waveform, sample_rate, target_sample_rate=16000):
+    #     if sample_rate != target_sample_rate:
+    #         waveform = torchaudio.functional.resample(waveform, sample_rate, target_sample_rate)
+    #     waveform = torch.mean(waveform, dim=0, keepdim=True)
+    #     return waveform