decoding utils update

utils.py

@@ -3,6 +3,7 @@ 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
@@ -10,29 +11,36 @@ import numpy as np
 from pydub import AudioSegment
 import shutil
 
+""" Diverse 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 index_has_substring(list, substring):
+    for i, s in enumerate(list):
+        if substring in s:
+            return i
+    return -1
 
-# Function to read from text from txt file:
-def readFromFile(path, isJSON=False):
-    with open(path, "r") as f:
-        if isJSON:
-            return json.load(f)
-        else:
-            return f.read()
+
+# 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:
@@ -42,21 +50,8 @@ def chain(input, funcs, *params):
     return res
 
 
-def to_beat_str(value, beat_res=8):
-    values = [
-        int(int(value * beat_res) / beat_res),
-        int(int(value * beat_res) % beat_res),
-        beat_res,
-    ]
-    return ".".join(map(str, values))
-
-
-def to_base10(beat_str):
-    integer, decimal, base = split_dots(beat_str)
-    return integer + decimal / base
-
-
 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(".")))
 
 
@@ -68,7 +63,39 @@ def get_datetime():
     return datetime.now().strftime("%Y%m%d_%H%M%S")
 
 
-def get_text(event):
+""" 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 "
@@ -77,13 +104,18 @@ def get_text(event):
         case "Track-End":
             return "TRACK_END "
         case "Instrument":
-            return f"INST={event.value} "
+            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_SHIFT={event.value} "
+            return f"TIME_DELTA={int_dec_base_to_delta(event.value, instrument)} "
         case "Note-On":
             return f"NOTE_ON={event.value} "
         case "Note-Off":
@@ -92,15 +124,66 @@ def get_text(event):
             return ""
 
 
-def get_event(text, value=None):
+""" 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 None
+            return Event("Track-Start", value)
         case "TRACK_END":
-            return None
+            return Event("Track-End", value)
         case "INST":
+            if value == "DRUMS":
+                value = "Drums"
             return Event("Instrument", value)
         case "BAR_START":
             return Event("Bar-Start", value)
@@ -109,7 +192,8 @@ def get_event(text, value=None):
         case "TIME_SHIFT":
             return Event("Time-Shift", value)
         case "TIME_DELTA":
-            return Event("Time-Shift", to_beat_str(int(value) / 4))
+            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":
@@ -118,39 +202,27 @@ def get_event(text, value=None):
             return None
 
 
-# TODO: Make this singleton
-def get_miditok():
-    pitch_range = range(0, 140)  # was (21, 109)
-    beat_res = {(0, 400): 8}
-    return MIDILike(pitch_range, beat_res)
-
+""" File utils"""
 
-class WriteTextMidiToFile:  # utils saving to file
-    def __init__(self, generate_midi, output_path):
-        self.generated_midi = generate_midi.generated_piece
-        self.output_path = output_path
-        self.hyperparameter_and_bars = generate_midi.piece_by_track
 
-    def hashing_seq(self):
-        self.current_time = get_datetime()
-        self.output_path_filename = f"{self.output_path}/{self.current_time}.json"
+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 wrapping_seq_hyperparameters_in_dict(self):
-        # assert type(self.generated_midi) is str, "error: generate_midi must be a string"
-        # assert (
-        #     type(self.hyperparameter_dict) is dict
-        # ), "error: feature_dict must be a dictionnary"
-        return {
-            "generate_midi": self.generated_midi,
-            "hyperparameters_and_bars": self.hyperparameter_and_bars,
-        }
 
-    def text_midi_to_file(self):
-        self.hashing_seq()
-        output_dict = self.wrapping_seq_hyperparameters_in_dict()
-        print(f"Token generate_midi written: {self.output_path_filename}")
-        writeToFile(self.output_path_filename, output_dict)
-        return self.output_path_filename
+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):
@@ -167,15 +239,33 @@ def get_files(directory, extension, recursive=False):
         return list(directory.glob(f"*.{extension}"))
 
 
-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
+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()
 
-    return wrapper
+
+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:
@@ -208,39 +298,3 @@ class FileCompressor:
         """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)
-
-
-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)
-
-
-def index_has_substring(list, substring):
-    for i, s in enumerate(list):
-        if substring in s:
-            return i
-    return -1