import functools import random import re import time import numpy as np import cvxpy as cp from bx.intervals.intersection import IntervalTree, Interval from data import VocabError from data_gen import Augmenter, AugmentError, RawExample, AugmentedExample from dataclasses import dataclass # TODO tests for this @dataclass class PitchIntervalTree: play_length: float pitch_intervals: dict[int, IntervalTree] pitch_interval_counts: dict[int, int] @classmethod def from_clip(cls, clip): play_length = 0 pitch_intervals = dict() pitch_interval_counts = dict() # TODO the intervals in this tree should be disjoint for n in clip: if n["offBeat"] <= n["onBeat"]: continue play_length += n["offBeat"] - n["onBeat"] ind = pitch_intervals.setdefault(n["note"], IntervalTree()) ind.insert_interval(Interval(n["onBeat"], n["offBeat"])) pitch_interval_counts[n["note"]] = ( pitch_interval_counts.get(n["note"], 0) + 1 ) return cls(play_length, pitch_intervals, pitch_interval_counts) def metric(self, other): min_play_length = min(self.play_length, other.play_length) if min_play_length == 0: return 0 overlap = 0 def visitor(tree, interval): nonlocal overlap for overlap_interval in tree.find(interval.start, interval.end): overlap += max( 0, min(interval.end, overlap_interval.end) - max(interval.start, overlap_interval.start), ) for pitch, intervals in self.pitch_intervals.items(): intervals_count = self.pitch_interval_counts[pitch] other_intervals = other.pitch_intervals.get(pitch) if other_intervals is None: continue other_intervals_count = other.pitch_interval_counts[pitch] smaller_tree = ( intervals if intervals_count < other_intervals_count else other_intervals ) larger_tree = other_intervals if smaller_tree is intervals else intervals smaller_tree.traverse(functools.partial(visitor, larger_tree)) return overlap / min_play_length # Augmenter for extracted-clips training data class ExtractedClipsAugmenter(Augmenter): def __init__( self, vocab, transpose_range=6, rebar_max_bars=4, rebar_prob=0.25, key_drop_prob=0.5, velocity_scale_prob=0.5, min_velocity=20, zero_accompany_prob=0.5, min_accompany=1, max_accompany=5, max_symbolic_length=2048, accompany_similar_prob=0.1, accompany_similar_threshold=0.95, max_qualified_accompany_to_main=3, cond_drop_prob=0.0, mask_main_prob=0.0, mask_main_max_bars=4, mask_accomp_prob=0.0, mask_accomp_max_bars=4, ): self.vocab = vocab self.transpose_range = transpose_range self.rebar_max_bars = rebar_max_bars self.rebar_prob = rebar_prob self.key_drop_prob = key_drop_prob self.velocity_scale_prob = velocity_scale_prob self.min_velocity = min_velocity self.zero_accompany_prob = zero_accompany_prob self.min_accompany = min_accompany self.max_accompany = max_accompany self.max_symbolic_length = max_symbolic_length self.accompany_similar_prob = accompany_similar_prob self.accompany_similar_threshold = accompany_similar_threshold self.max_qualified_accompany_to_main = max_qualified_accompany_to_main self.cond_drop_prob = cond_drop_prob self.mask_main_prob = mask_main_prob self.mask_main_max_bars = mask_main_max_bars self.mask_accomp_prob = mask_accomp_prob self.mask_accomp_max_bars = mask_accomp_max_bars @classmethod def from_config(cls, vocab, common_config, my_config): return cls( vocab, int(my_config["transpose_range"]), int(my_config["rebar_max_bars"]), float(my_config["rebar_prob"]), float(my_config["key_drop_prob"]), float(my_config["velocity_scale_prob"]), int(my_config["min_velocity"]), float(my_config["zero_accompany_prob"]), int(my_config["min_accompany"]), int(my_config["max_accompany"]), int(common_config["seq_len_max"]), float(my_config["accompany_similar_prob"]), float(my_config["accompany_similar_threshold"]), int(my_config["max_qualified_accompany_to_main"]), float(my_config["cond_drop_prob"]), float(my_config["mask_main_prob"]), float(my_config["mask_main_max_bars"]), float(my_config["mask_accomp_prob"]), float(my_config["mask_accomp_max_bars"]), ) # TODO this really needs to produce better results! def _augment_descs(self, descs): if random.random() < self.cond_drop_prob: return "" def _fixup_desc(desc): spaced = re.sub(r"[\r\n_/-]", " ", desc) demidied = re.sub(r"\.midi?$", "", spaced) stripped = re.sub( r"^(Artist|Title|Composer|Lyricist|Subtitle|Scriptures|Editor|Source|Notes|Arr)\s*[:.]?\s*", "", demidied.strip(), ) return " ".join(tok.lower() for tok in stripped.split(" ") if tok) tags_prefix = "" descs_by_type = {} for desc in descs: descs_by_type.setdefault(desc["type"], set()).add( _fixup_desc(desc["description"]) ) def take(type, min_count): xs = list(descs_by_type.get(type, [])) random.shuffle(xs) return xs[: max(min_count, random.randint(0, len(xs)))] descs_out = ( take("artist_or_title", 1) + take("musescore_text", 2) + take("imslp_text", 2) + take("file_name", 1) + take("tag_grab_bag", 0) ) descs_out = list(set(descs_out)) random.shuffle(descs_out) return tags_prefix + " ".join(descs_out) def _augment_clips(self, examples, key): if key is not None and random.random() < self.key_drop_prob: key = None allnotes = [n for example in examples for n in example] if len(allnotes) == 0: return examples, key if any(n["note"] < 1000 for n in allnotes): highest_note = max(n["note"] for n in allnotes if n["note"] < 1000) lowest_note = min(n["note"] for n in allnotes if n["note"] < 1000) headroom_high = self.vocab.pitch_max - highest_note headroom_low = lowest_note - self.vocab.pitch_min headroom = headroom_high + headroom_low if ( headroom < 0 or -headroom_high > self.transpose_range or -headroom_low > self.transpose_range ): raise AugmentError("Not enough headroom to transpose") transpose = random.randint( -min(headroom_low, self.transpose_range), min(headroom_high, self.transpose_range), ) examples = [ [ dict(n, note=n["note"] + (transpose if n["note"] < 1000 else 0)) for n in example ] for example in examples ] key = None if key is None else (key + transpose) % 12 + key // 12 if random.random() < self.rebar_prob: latest_event_time = max(n["offBeat"] for n in allnotes) shift = 4 * random.randint( 0, max(0, self.rebar_max_bars - latest_event_time // 4) ) examples = [ [ dict(n, onBeat=n["onBeat"] + shift, offBeat=n["offBeat"] + shift) for n in example ] for example in examples ] if random.random() < self.velocity_scale_prob: least_velocity = max(1, min(n["onVelocity"] for n in allnotes)) greatest_velocity = max(1, max(n["onVelocity"] for n in allnotes)) scale_min = self.min_velocity / least_velocity scale_max = 127 / greatest_velocity if scale_max < scale_min: raise AugmentError("Not enough velocity range to scale") scale = random.uniform(scale_min, scale_max) examples = [ [ dict(n, onVelocity=int(round(n["onVelocity"] * scale))) for n in example ] for example in examples ] return examples, key def _bin_pack( self, main_example, qual_accompaniments, all_accompaniments, target_num_extra_main_clips, target_num_accompaniments, ): all_clips = [main_example] # allow VocabError to escape for the main example sym_lengths = [ self.vocab.fast_estimate_length(main_example, include_header=False) ] for accomp in all_accompaniments: try: sym_lengths.append( self.vocab.fast_estimate_length(accomp, include_header=False) ) all_clips.append(accomp) except VocabError: continue all_accompaniments = all_clips[1:] sym_lengths = np.array(sym_lengths) allowed_in_main = np.array( [1] # <- main_example + [ any(accomp is qual_accomp for qual_accomp in qual_accompaniments) for accomp in all_accompaniments ] ) allowed_in_main_idxs = np.nonzero(allowed_in_main)[0] is_in_main = cp.Variable(len(all_clips), boolean=True) is_in_accomp = cp.Variable(len(all_clips), boolean=True) similarity_constraints = [] if random.random() >= self.accompany_similar_prob: all_pi_trees = [PitchIntervalTree.from_clip(clip) for clip in all_clips] for i in allowed_in_main_idxs: for j in range(len(all_clips)): if i != j: similarity = all_pi_trees[i].metric(all_pi_trees[j]) if similarity > self.accompany_similar_threshold: similarity_constraints.append( is_in_main[i] + is_in_accomp[j] <= 1 ) # multiply decision variables by these weights to choose a random solution among the optimal ones perturb = 1.0 + 0.1 * np.random.uniform(0, 1, len(all_clips)) / len(all_clips) objective = cp.Maximize((is_in_main + is_in_accomp) @ perturb) constraints = [ is_in_main[0] == 1, is_in_main + is_in_accomp <= 1, cp.sum(is_in_main) <= 1 + target_num_extra_main_clips, cp.sum(is_in_accomp) <= target_num_accompaniments, is_in_main <= allowed_in_main, sym_lengths @ (is_in_main + is_in_accomp) <= self.max_symbolic_length - 4, ] problem = cp.Problem(objective, constraints + similarity_constraints) t0 = time.time() problem.solve() dt = time.time() - t0 if objective.value is None and len(similarity_constraints) > 0: problem = cp.Problem(objective, constraints) t0 = time.time() problem.solve() dt = time.time() - t0 if objective.value is None: return main_example, [] main_comps = [ all_clips[i] for i in range(len(all_clips)) if is_in_main.value[i] ] accomp_comps = [ all_clips[i] for i in range(len(all_clips)) if is_in_accomp.value[i] ] return ( [n for comp in main_comps for n in comp], [n for comp in accomp_comps for n in comp], ) def _mask_clip(self, clip, max_bars): mask_start = random.random() * max_bars * 4 return [n for n in clip if n["onBeat"] >= mask_start] def _mask_clip_right(self, clip): if len(clip) == 0: return clip earliest_evt = min(n["onBeat"] for n in clip) latest_evt = max(n["offBeat"] for n in clip) mask_start = earliest_evt + random.random() * (latest_evt - earliest_evt) return [n for n in clip if n["offBeat"] <= mask_start] def augment(self, example: RawExample) -> AugmentedExample: aug_qual_accompaniments = [ accomp for accomp, accomp_tags in zip(example.accomps, example.accomp_tagss) if "Trainable as Main" in accomp_tags ] target_num_extra_main_clips = random.randint( 0, self.max_qualified_accompany_to_main ) target_num_accomps = ( 0 if random.random() < self.zero_accompany_prob else random.randint(self.min_accompany, self.max_accompany) ) composite_main, composite_accomp = self._bin_pack( example.example, aug_qual_accompaniments, example.accomps, target_num_extra_main_clips, target_num_accomps, ) [aug_main, aug_accomp], _ = self._augment_clips( [composite_main, composite_accomp], None ) if random.random() < self.mask_accomp_prob: aug_accomp = self._mask_clip(aug_accomp, self.mask_accomp_max_bars) if random.random() < self.mask_accomp_prob: aug_accomp = self._mask_clip_right(aug_accomp) if len(aug_accomp) > 0 and random.random() < self.mask_main_prob: aug_main = self._mask_clip(aug_main, self.mask_main_max_bars) return AugmentedExample( id=example.id, desc=self._augment_descs(example.descs), example=aug_main, accomp=aug_accomp, )