InfiniteMusic/data.py
2025-10-15 02:19:33 -04:00

74 lines
1.8 KiB
Python

import librosa
import numpy as np
from pathlib import Path
from multiprocessing import Pool, cpu_count
import mlflow
SAMPLE_RATE = 22050
#@mlflow.trace
def process_file(file_path):
"""
Load 10 second chunks single song.
"""
y, sr = librosa.load(file_path, mono=True, sr=SAMPLE_RATE)
if(not sr == SAMPLE_RATE):
return []
size = int(SAMPLE_RATE * 10)
sample_len = len(y)
file_chunks = []
for start_pos in range(0, sample_len, size):
end = start_pos + size
if end <= sample_len:
chunk = y[start_pos:end]
#chunk = librosa.feature.melspectrogram(y=chunk,sr=SAMPLE_RATE)
#chunk = ((librosa.amplitude_to_db(chunk,ref=np.max)+40)/40)
file_chunks.append(chunk)
return file_chunks
#@mlflow.trace
def load():
"""
Load 10 second chunks of songs.
"""
audio = []
files = list(Path("./data/").glob("*.mp3"))
#files = files[:12]
with Pool(cpu_count()) as pool:
chunk_list = pool.map(process_file, files)
for l in chunk_list:
audio.extend(l)
return audio
##DEP
def audio_split(audio):
"""
Split 10 seconds of audio to 2 5 second clips
"""
size = int(SAMPLE_RATE*5)
x = audio[:size]
y = audio[size:size*2]
x = librosa.feature.melspectrogram(y=x, sr=SAMPLE_RATE)
y = librosa.feature.melspectrogram(y=y, sr=SAMPLE_RATE)
x = ((librosa.amplitude_to_db(x,ref=np.max)+80)/80)
y = ((librosa.amplitude_to_db(y,ref=np.max)+80)/80)
return x,y
def dataset(chunks):
"""
convert 10 second chunks to dataset
"""
x,y=[],[]
with Pool(cpu_count()) as pool:
audio_list = pool.map(audio_split,chunks)
for (ax,ay) in audio_list:
x.append(ax)
y.append(ay)
return x,y