Update audio2text/a2t.py

audio2text/a2t.py

@@ -1,9 +1,10 @@
+import librosa
 import numpy as np
-
 from .init import pipe
 
 TASK = "transcribe"
 BATCH_SIZE = 16
+LIMIT = 60
 
 class A2T:
     def __init__(self, mic):
@@ -13,12 +14,17 @@ class A2T:
         if inputs is  None:
             print("Inputs None")
 
-        transcribed_text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task, "language": "english"})
-        print("transcribed_text : ", transcribed_text)
-        return transcribed_text["text"]
+        transcribed_text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task, "language": "english"})["text"]
+        return transcribed_text
 
-    def __preprocces(self, raw: np.ndarray):
+    def __preprocces(self, raw: np.ndarray, sampling_rate: int):
         chunk = raw.astype(np.float32) / 32768.0
+
+        if sampling_rate > 16000:
+            chunk = librosa.resample(chunk, orig_sr=sampling_rate, target_sr=16000)
+
+        chunk = chunk[:16000*LIMIT]
+                
         return chunk
         
     def predict(self):
@@ -26,7 +32,8 @@ class A2T:
             if self.mic is not None:
                 chunk = self.mic.get_array_of_samples()
                 chunk = np.array(chunk, dtype=np.int16)
-                audio = self.__preprocces(chunk)
+                sampling_rate = self.mic.frame_rate
+                audio = self.__preprocces(raw=chunk, sampling_rate=sampling_rate)
                 print(f"audio : {audio} \n shape : {audio.shape} \n max : {np.max(audio)}")
             else:
                 raise Exception("please provide audio")