conversion.py

from riffusion_helper.spectrogram_image_converter import SpectrogramImageConverter
from riffusion_helper.spectrogram_params import SpectrogramParams
import pydub
from PIL import Image

params = SpectrogramParams()

converter = SpectrogramImageConverter(params=params)

def save_spectrogram_image_from_audio(filename: str):
    audio = pydub.AudioSegment.from_file(f'{filename}')

    print(len(audio))

    name = filename.split('/')[-1].split('.')[0]

    for i in range(0, len(audio), 5000):
        clip = audio[i:i+5000]
        image = converter.spectrogram_image_from_audio(clip)
        image.save(f'./dataset/spectrograms/{name}_{i}.png')
        print(f'./dataset/spectrograms/{name}_{i}.png')
    
    return True

def create_audio_from_spectrogram_image(filename: str):
    image = Image.open(filename)
    audio = converter.audio_from_spectrogram_image(image)
    audio.export(f'./{filename.split("/")[-1].split(".")[0]}.mp3', format="mp3")

create_audio_from_spectrogram_image("soothingSlowMusic.jpg")