testtranscription.py
1 """ 2 Transcription module tests 3 """ 4 5 import unittest 6 7 import numpy as np 8 import soundfile as sf 9 10 from scipy import signal 11 12 from txtai.pipeline import Transcription 13 14 # pylint: disable=C0411 15 from utils import Utils 16 17 18 class TestTranscription(unittest.TestCase): 19 """ 20 Transcription tests. 21 """ 22 23 def testArray(self): 24 """ 25 Test audio data to text transcription 26 """ 27 28 transcribe = Transcription() 29 30 # Read audio data 31 raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav") 32 33 self.assertEqual(transcribe((raw, samplerate)), "Make huge profits without working make up to one hundred thousand dollars a day") 34 self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day") 35 36 def testChunks(self): 37 """ 38 Test splitting transcription into chunks 39 """ 40 41 transcribe = Transcription() 42 43 result = transcribe(Utils.PATH + "/Make_huge_profits.wav", join=False)[0] 44 45 self.assertIsInstance(result["raw"], np.ndarray) 46 self.assertIsNotNone(result["rate"]) 47 self.assertEqual(result["text"], "Make huge profits without working make up to one hundred thousand dollars a day") 48 49 def testFile(self): 50 """ 51 Test audio file to text transcription 52 """ 53 54 transcribe = Transcription() 55 56 self.assertEqual( 57 transcribe(Utils.PATH + "/Make_huge_profits.wav"), "Make huge profits without working make up to one hundred thousand dollars a day" 58 ) 59 60 def testGenerateArguments(self): 61 """ 62 Test transcription with generation keyword arguments 63 """ 64 65 transcribe = Transcription() 66 67 # Read audio data 68 raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav") 69 70 self.assertEqual( 71 transcribe(raw, samplerate, language="English", task="transcribe"), 72 "Make huge profits without working make up to one hundred thousand dollars a day", 73 ) 74 75 def testResample(self): 76 """ 77 Test resampled audio file to text transcription 78 """ 79 80 transcribe = Transcription() 81 82 # Read audio data 83 raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav") 84 85 # Resample for testing 86 samples = round(len(raw) * float(22050) / samplerate) 87 raw, samplerate = signal.resample(raw, samples), 22050 88 89 self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day") 90 91 def testStereo(self): 92 """ 93 Test audio file in stereo to text transcription 94 """ 95 96 transcribe = Transcription() 97 98 # Read audio data 99 raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav") 100 101 # Convert mono to stereo 102 raw = np.column_stack((raw, raw)) 103 104 self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")