implement spectrogram method

app.py

@@ -58,6 +58,26 @@ def check_outputs_folder(folder_path):
     else:
         print(f'The folder {folder_path} does not exist.')
 
+def plot_spectrogram(wav_file, output_image):
+    # Read the WAV file
+    sample_rate, audio_data = wavfile.read(wav_file)
+
+    # Check if audio_data is stereo (2 channels) and convert it to mono (1 channel) if needed
+    if len(audio_data.shape) == 2:
+        audio_data = audio_data.mean(axis=1)
+
+    # Create a plot for the spectrogram
+    plt.figure(figsize=(10, 4))
+    plt.specgram(audio_data, Fs=sample_rate, NFFT=1024, noverlap=512, cmap='inferno', aspect='auto')
+    plt.title('Spectrogram')
+    plt.xlabel('Time [s]')
+    plt.ylabel('Frequency [Hz]')
+
+    # Save the plot as an image file
+    plt.colorbar(label='Intensity [dB]')
+    plt.savefig(output_image)
+    plt.close()
+
 def infer(video_in):
 
     # check if 'outputs' dir exists and empty it if necessary
@@ -92,7 +112,9 @@ def infer(video_in):
     print_directory_contents('./outputs/tmp')
     wave_files = get_wav_files('./outputs/tmp')
     print(wave_files)
-    return wave_files[0]
+    plot_spectrogram(wave_files[0], 'spectrogram.png')
+    
+    return wave_files[0], 'spectrogram.png'
 
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
@@ -100,11 +122,12 @@ with gr.Blocks() as demo:
         video_in = gr.Video(label='Video IN')
         submit_btn = gr.Button("Submit")
         output_sound = gr.Audio(label="Audio OUT")
+        output_spectrogram = gr.Image(label='Spectrogram')
         #output_sound = gr.Textbox(label="Audio OUT")
     submit_btn.click(
         fn = infer,
         inputs = [video_in],
-        outputs = [output_sound],
+        outputs = [output_sound, output_spectrogram],
         show_api = False
     )
 demo.launch(show_api=False, show_error=True)