from datetime import datetime
from miditok import Event, MIDILike
import os
import json
from time import perf_counter
from constants import DRUMS_BEAT_QUANTIZATION, NONE_DRUMS_BEAT_QUANTIZATION
from joblib import Parallel, delayed
from zipfile import ZipFile, ZIP_DEFLATED
from scipy.io.wavfile import write
import numpy as np
from pydub import AudioSegment
import shutil
""" Diverse utils"""
def index_has_substring(list, substring):
for i, s in enumerate(list):
if substring in s:
return i
return -1
# TODO: Make this singleton
def get_miditok():
pitch_range = range(0, 127) # was (21, 109)
beat_res = {(0, 400): 8}
return MIDILike(pitch_range, beat_res)
def timeit(func):
def wrapper(*args, **kwargs):
start = perf_counter()
result = func(*args, **kwargs)
end = perf_counter()
print(f"{func.__name__} took {end - start:.2f} seconds to run.")
return result
return wrapper
def chain(input, funcs, *params):
"""Chain functions together, passing the output of one function as the input of the next."""
res = input
for func in funcs:
try:
res = func(res, *params)
except TypeError:
res = func(res)
return res
def split_dots(value):
"""Splits a string separated by dots "a.b.c" into a list of integers [a, b, c]"""
return list(map(int, value.split(".")))
def compute_list_average(l):
return sum(l) / len(l)
def get_datetime():
return datetime.now().strftime("%Y%m%d_%H%M%S")
""" Encoding functions """
def int_dec_base_to_beat(beat_str):
"""
Converts "integer.decimal.base" (str, from miditok) into beats
e.g. "0.4.8" = 0 + 4/8 = 0.5
Args:
- beat_str: "integer.decimal.base"
Returns:
- beats: float
"""
integer, decimal, base = split_dots(beat_str)
return integer + decimal / base
def int_dec_base_to_delta(beat_str, instrument="drums"):
"""converts the time shift to time_delta according to Tristan's encoding scheme
Drums TIME_DELTA are quantized according to DRUMS_BEAT_QUANTIZATION
Other Instrument TIME_DELTA are quantized according to NONE_DRUMS_BEAT_QUANTIZATION
"""
beat_res = (
DRUMS_BEAT_QUANTIZATION
if instrument.lower() == "drums"
else NONE_DRUMS_BEAT_QUANTIZATION
)
time_delta = int_dec_base_to_beat(beat_str) * beat_res
return time_delta.__int__()
def get_text(event, instrument="drums"):
"""Converts an event into a string for the midi-text format"""
match event.type:
case "Piece-Start":
return "PIECE_START "
case "Track-Start":
return "TRACK_START "
case "Track-End":
return "TRACK_END "
case "Instrument":
if str(event.value).lower() == "drums":
return f"INST=DRUMS "
else:
return f"INST={event.value} "
case "Density":
return f"DENSITY={event.value} "
case "Bar-Start":
return "BAR_START "
case "Bar-End":
return "BAR_END "
case "Time-Shift":
return f"TIME_DELTA={int_dec_base_to_delta(event.value, instrument)} "
case "Note-On":
return f"NOTE_ON={event.value} "
case "Note-Off":
return f"NOTE_OFF={event.value} "
case _:
return ""
""" Decoding functions """
def time_delta_to_beat(time_delta, instrument="drums"):
"""
Converts TIME_DELTA (from midi-text) to beats according to Tristan's encoding scheme
Args:
- time_delta: int (TIME_DELTA)
- instrument: str ("Drums" or other instrument): used to determine the quantization resolution defined on constants.py
Returns:
- beats: float
"""
beat_res = (
DRUMS_BEAT_QUANTIZATION
if instrument.lower() == "drums"
else NONE_DRUMS_BEAT_QUANTIZATION
)
beats = float(time_delta) / beat_res
return beats
def beat_to_int_dec_base(beat, beat_res=8):
"""
Converts beats into "integer.decimal.base" (str) for miditok
Args:
- beat_str: "integer.decimal.base"
Returns:
- beats: float (e.g. "0.4.8" = 0 + 4/8 = 0.5)
"""
int_dec_base = [
int((beat * beat_res) // beat_res),
int((beat * beat_res) % beat_res),
beat_res,
]
return ".".join(map(str, int_dec_base))
def time_delta_to_int_dec_base(time_delta, instrument="drums"):
return chain(
time_delta,
[
time_delta_to_beat,
beat_to_int_dec_base,
],
instrument,
)
def get_event(text, value=None, instrument="drums"):
"""Converts a midi-text like event into a miditok like event"""
match text:
case "PIECE_START":
return Event("Piece-Start", value)
case "TRACK_START":
return Event("Track-Start", value)
case "TRACK_END":
return Event("Track-End", value)
case "INST":
if value == "DRUMS":
value = "Drums"
return Event("Instrument", value)
case "BAR_START":
return Event("Bar-Start", value)
case "BAR_END":
return Event("Bar-End", value)
case "TIME_SHIFT":
return Event("Time-Shift", value)
case "TIME_DELTA":
return Event("Time-Shift", time_delta_to_int_dec_base(value, instrument))
# return Event("Time-Shift", to_beat_str(int(value) / 4))
case "NOTE_ON":
return Event("Note-On", value)
case "NOTE_OFF":
return Event("Note-Off", value)
case _:
return None
""" File utils"""
def writeToFile(path, content):
if type(content) is dict:
with open(f"{path}", "w") as json_file:
json.dump(content, json_file)
else:
if type(content) is not str:
content = str(content)
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
f.write(content)
def readFromFile(path, isJSON=False):
with open(path, "r") as f:
if isJSON:
return json.load(f)
else:
return f.read()
def get_files(directory, extension, recursive=False):
"""
Given a directory, get a list of the file paths of all files matching the
specified file extension.
directory: the directory to search as a Path object
extension: the file extension to match as a string
recursive: whether to search recursively in the directory or not
"""
if recursive:
return list(directory.rglob(f"*.{extension}"))
else:
return list(directory.glob(f"*.{extension}"))
def load_jsonl(filepath):
"""Load a jsonl file"""
with open(filepath, "r") as f:
data = [json.loads(line) for line in f]
return data
def write_mp3(waveform, output_path, bitrate="92k"):
"""
Write a waveform to an mp3 file.
output_path: Path object for the output mp3 file
waveform: numpy array of the waveform
bitrate: bitrate of the mp3 file (64k, 92k, 128k, 256k, 312k)
"""
# write the wav file
wav_path = output_path.with_suffix(".wav")
write(wav_path, 44100, waveform.astype(np.float32))
# compress the wav file as mp3
AudioSegment.from_wav(wav_path).export(output_path, format="mp3", bitrate=bitrate)
# remove the wav file
wav_path.unlink()
def copy_file(input_file, output_dir):
"""Copy an input file to the output_dir"""
output_file = output_dir / input_file.name
shutil.copy(input_file, output_file)
class FileCompressor:
def __init__(self, input_directory, output_directory, n_jobs=-1):
self.input_directory = input_directory
self.output_directory = output_directory
self.n_jobs = n_jobs
# File compression and decompression
def unzip_file(self, file):
"""uncompress single zip file"""
with ZipFile(file, "r") as zip_ref:
zip_ref.extractall(self.output_directory)
def zip_file(self, file):
"""compress a single text file to a new zip file and delete the original"""
output_file = self.output_directory / (file.stem + ".zip")
with ZipFile(output_file, "w") as zip_ref:
zip_ref.write(file, arcname=file.name, compress_type=ZIP_DEFLATED)
file.unlink()
@timeit
def unzip(self):
"""uncompress all zip files in folder"""
files = get_files(self.input_directory, extension="zip")
Parallel(n_jobs=self.n_jobs)(delayed(self.unzip_file)(file) for file in files)
@timeit
def zip(self):
"""compress all text files in folder to new zip files and remove the text files"""
files = get_files(self.output_directory, extension="txt")
Parallel(n_jobs=self.n_jobs)(delayed(self.zip_file)(file) for file in files)