diff --git a/MANIFEST.in b/MANIFEST.in index b2858a4..6e2c0f4 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,6 +1,6 @@ include *.txt tox.ini *.rst *.md LICENSE include catalog-info.yaml include Dockerfile .dockerignore -recursive-include tests *.py *.wav *.npz *.jams *.zip +recursive-include tests *.py *.wav *.npz *.jams *.zip *.midi *.csv *.json recursive-include basic_pitch *.py *.md recursive-include basic_pitch/saved_models *.index *.pb variables.data* *.mlmodel *.json *.onnx *.tflite *.bin diff --git a/basic_pitch/data/datasets/maestro.py b/basic_pitch/data/datasets/maestro.py new file mode 100644 index 0000000..e4ab3e7 --- /dev/null +++ b/basic_pitch/data/datasets/maestro.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python +# encoding: utf-8 +# +# Copyright 2024 Spotify AB +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import logging +import os +import sys +import tempfile +import time +from typing import Any, Dict, List, TextIO, Tuple + +import apache_beam as beam +import mirdata + +from basic_pitch.data import commandline, pipeline + + +def read_in_chunks(file_object: TextIO, chunk_size: int = 1024) -> Any: + """Lazy function (generator) to read a file piece by piece. + Default chunk size: 1k.""" + while True: + data = file_object.read(chunk_size) + if not data: + break + yield data + + +class MaestroInvalidTracks(beam.DoFn): + DOWNLOAD_ATTRIBUTES = ["audio_path"] + + def __init__(self, source: str) -> None: + self.source = source + + def setup(self) -> None: + # Oddly enough we dont want to include the gcs bucket uri. + # Just the path within the bucket + self.maestro_remote = mirdata.initialize("maestro", data_home=self.source) + self.filesystem = beam.io.filesystems.FileSystems() + + def process(self, element: Tuple[str, str], *args: Tuple[Any, Any], **kwargs: Dict[str, Any]) -> Any: + import tempfile + import sox + + track_id, split = element + logging.info(f"Processing (track_id, split): ({track_id}, {split})") + + track_remote = self.maestro_remote.track(track_id) + with tempfile.TemporaryDirectory() as local_tmp_dir: + maestro_local = mirdata.initialize("maestro", local_tmp_dir) + track_local = maestro_local.track(track_id) + + for attribute in self.DOWNLOAD_ATTRIBUTES: + source = getattr(track_remote, attribute) + destination = getattr(track_local, attribute) + os.makedirs(os.path.dirname(destination), exist_ok=True) + with self.filesystem.open(source) as s, open(destination, "wb") as d: + for piece in read_in_chunks(s): + d.write(piece) + + # 15 minutes * 60 seconds/minute + if sox.file_info.duration(track_local.audio_path) >= 15 * 60: + return None + + yield beam.pvalue.TaggedOutput(split, track_id) + + +class MaestroToTfExample(beam.DoFn): + DOWNLOAD_ATTRIBUTES = ["audio_path", "midi_path"] + + def __init__(self, source: str, download: bool): + self.source = source + self.download = download + + def setup(self) -> None: + import apache_beam as beam + import mirdata + + # Oddly enough we dont want to include the gcs bucket uri. + # Just the path within the bucket + self.maestro_remote = mirdata.initialize("maestro", data_home=self.source) + self.filesystem = beam.io.filesystems.FileSystems() + if self.download: + self.maestro_remote.download() + + def process(self, element: List[str], *args: Tuple[Any, Any], **kwargs: Dict[str, Any]) -> List[Any]: + import tempfile + + import numpy as np + import sox + + from basic_pitch.constants import ( + AUDIO_N_CHANNELS, + AUDIO_SAMPLE_RATE, + FREQ_BINS_CONTOURS, + FREQ_BINS_NOTES, + ANNOTATION_HOP, + N_FREQ_BINS_NOTES, + N_FREQ_BINS_CONTOURS, + ) + from basic_pitch.data import tf_example_serialization + + logging.info(f"Processing {element}") + batch = [] + + for track_id in element: + track_remote = self.maestro_remote.track(track_id) + with tempfile.TemporaryDirectory() as local_tmp_dir: + maestro_local = mirdata.initialize("maestro", local_tmp_dir) + track_local = maestro_local.track(track_id) + + for attribute in self.DOWNLOAD_ATTRIBUTES: + source = getattr(track_remote, attribute) + destination = getattr(track_local, attribute) + os.makedirs(os.path.dirname(destination), exist_ok=True) + with self.filesystem.open(source) as s, open(destination, "wb") as d: + # d.write(s.read()) + for piece in read_in_chunks(s): + d.write(piece) + + local_wav_path = f"{track_local.audio_path}_tmp.wav" + + tfm = sox.Transformer() + tfm.rate(AUDIO_SAMPLE_RATE) + tfm.channels(AUDIO_N_CHANNELS) + tfm.build(track_local.audio_path, local_wav_path) + + duration = sox.file_info.duration(local_wav_path) + time_scale = np.arange(0, duration + ANNOTATION_HOP, ANNOTATION_HOP) + n_time_frames = len(time_scale) + + note_indices, note_values = track_local.notes.to_sparse_index(time_scale, "s", FREQ_BINS_NOTES, "hz") + onset_indices, onset_values = track_local.notes.to_sparse_index( + time_scale, "s", FREQ_BINS_NOTES, "hz", onsets_only=True + ) + contour_indices, contour_values = track_local.notes.to_sparse_index( + time_scale, "s", FREQ_BINS_CONTOURS, "hz" + ) + + batch.append( + tf_example_serialization.to_transcription_tfexample( + track_local.track_id, + "maestro", + local_wav_path, + note_indices, + note_values, + onset_indices, + onset_values, + contour_indices, + contour_values, + (n_time_frames, N_FREQ_BINS_NOTES), + (n_time_frames, N_FREQ_BINS_CONTOURS), + ) + ) + return [batch] + + +def create_input_data(source: str) -> List[Tuple[str, str]]: + import apache_beam as beam + + filesystem = beam.io.filesystems.FileSystems() + + with tempfile.TemporaryDirectory() as tmpdir: + maestro = mirdata.initialize("maestro", data_home=tmpdir) + metadata_path = maestro._index["metadata"]["maestro-v2.0.0"][0] + with filesystem.open( + os.path.join(source, metadata_path), + ) as s, open(os.path.join(tmpdir, metadata_path), "wb") as d: + d.write(s.read()) + + return [(track_id, track.split) for track_id, track in maestro.load_tracks().items()] + + +def main(known_args: argparse.Namespace, pipeline_args: List[str]) -> None: + time_created = int(time.time()) + destination = commandline.resolve_destination(known_args, time_created) + + # TODO: Remove or abstract for foss + pipeline_options = { + "runner": known_args.runner, + "job_name": f"maestro-tfrecords-{time_created}", + "machine_type": "e2-highmem-4", + "num_workers": 25, + "disk_size_gb": 128, + "experiments": ["use_runner_v2", "no_use_multiple_sdk_containers"], + "save_main_session": True, + "sdk_container_image": known_args.sdk_container_image, + "job_endpoint": known_args.job_endpoint, + "environment_type": "DOCKER", + "environment_config": known_args.sdk_container_image, + } + input_data = create_input_data(known_args.source) + pipeline.run( + pipeline_options, + pipeline_args, + input_data, + MaestroToTfExample(known_args.source, download=True), + MaestroInvalidTracks(known_args.source), + destination, + known_args.batch_size, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + commandline.add_default(parser, os.path.basename(os.path.splitext(__file__)[0])) + commandline.add_split(parser) + known_args, pipeline_args = parser.parse_known_args(sys.argv) + + main(known_args, pipeline_args) diff --git a/basic_pitch/data/download.py b/basic_pitch/data/download.py index 085ae75..e11229a 100644 --- a/basic_pitch/data/download.py +++ b/basic_pitch/data/download.py @@ -20,12 +20,19 @@ from basic_pitch.data import commandline from basic_pitch.data.datasets.guitarset import main as guitarset_main from basic_pitch.data.datasets.ikala import main as ikala_main +from basic_pitch.data.datasets.maestro import main as maestro_main from basic_pitch.data.datasets.medleydb_pitch import main as medleydb_pitch_main logger = logging.getLogger() logger.setLevel(logging.INFO) -DATASET_DICT = {"guitarset": guitarset_main, "ikala": ikala_main, "medleydb_pitch": medleydb_pitch_main} + +DATASET_DICT = { + "guitarset": guitarset_main, + "ikala": ikala_main, + "maestro": maestro_main, + "medleydb_pitch": medleydb_pitch_main, +} def main() -> None: diff --git a/pyproject.toml b/pyproject.toml index fa7001b..4e23b0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,8 @@ test = [ "coverage>=5.0.2", "pytest>=6.1.1", "pytest-mock", + "wave", + "mido" ] tf = [ "tensorflow>=2.4.1,<2.15.1; platform_system != 'Darwin'", diff --git a/tests/data/test_guitarset.py b/tests/data/test_guitarset.py index 2ff8ebd..708f59d 100644 --- a/tests/data/test_guitarset.py +++ b/tests/data/test_guitarset.py @@ -33,7 +33,7 @@ TRACK_ID = "00_BN1-129-Eb_comp" -def test_guitar_set_to_tf_example(tmpdir: str) -> None: +def test_guitarset_to_tf_example(tmpdir: str) -> None: input_data: List[str] = [TRACK_ID] with TestPipeline() as p: ( @@ -51,7 +51,7 @@ def test_guitar_set_to_tf_example(tmpdir: str) -> None: assert len(data) != 0 -def test_guitar_set_invalid_tracks(tmpdir: str) -> None: +def test_guitarset_invalid_tracks(tmpdir: str) -> None: split_labels = ["train", "test", "validation"] input_data = [(str(i), split) for i, split in enumerate(split_labels)] with TestPipeline() as p: @@ -73,7 +73,7 @@ def test_guitar_set_invalid_tracks(tmpdir: str) -> None: assert fp.read().strip() == str(i) -def test_create_input_data() -> None: +def test_guitarset_create_input_data() -> None: data = create_input_data(train_percent=0.33, validation_percent=0.33) data.sort(key=lambda el: el[1]) # sort by split tolerance = 0.1 @@ -81,7 +81,7 @@ def test_create_input_data() -> None: assert (0.33 - tolerance) * len(data) <= len(list(group)) <= (0.33 + tolerance) * len(data) -def test_create_input_data_overallocate() -> None: +def test_guitarset_create_input_data_overallocate() -> None: try: create_input_data(train_percent=0.6, validation_percent=0.6) except AssertionError: diff --git a/tests/data/test_ikala.py b/tests/data/test_ikala.py index 09682a4..66756cc 100644 --- a/tests/data/test_ikala.py +++ b/tests/data/test_ikala.py @@ -51,15 +51,15 @@ def test_ikala_invalid_tracks(tmpdir: str) -> None: assert fp.read().strip() == str(i) -def test_create_input_data() -> None: +def test_ikala_create_input_data() -> None: data = create_input_data(train_percent=0.5) data.sort(key=lambda el: el[1]) # sort by split - tolerance = 0.05 - for key, group in itertools.groupby(data, lambda el: el[1]): + tolerance = 0.1 + for _, group in itertools.groupby(data, lambda el: el[1]): assert (0.5 - tolerance) * len(data) <= len(list(group)) <= (0.5 + tolerance) * len(data) -def test_create_input_data_overallocate() -> None: +def test_ikala_create_input_data_overallocate() -> None: try: create_input_data(train_percent=1.1) except AssertionError: diff --git a/tests/data/test_maestro.py b/tests/data/test_maestro.py new file mode 100644 index 0000000..e910963 --- /dev/null +++ b/tests/data/test_maestro.py @@ -0,0 +1,192 @@ +#!/usr/bin/env python +# encoding: utf-8 +# +# Copyright 2024 Spotify AB +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import logging +import numpy as np +import os +import pathlib +import wave + +from mido import MidiFile, MidiTrack, Message +from typing import List + +import apache_beam as beam +from apache_beam.testing.test_pipeline import TestPipeline + +from basic_pitch.data.datasets.maestro import ( + MaestroToTfExample, + MaestroInvalidTracks, + create_input_data, +) +from basic_pitch.data.pipeline import WriteBatchToTfRecord + +RESOURCES_PATH = pathlib.Path(__file__).parent.parent / "resources" +MAESTRO_TEST_DATA_PATH = RESOURCES_PATH / "data" / "maestro" + +TRAIN_TRACK_ID = "2004/MIDI-Unprocessed_SMF_05_R1_2004_01_ORIG_MID--AUDIO_05_R1_2004_03_Track03_wav" +VALID_TRACK_ID = "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_06_Track06_wav" +TEST_TRACK_ID = "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_08_Track08_wav" +GT_15M_TRACK_ID = "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_05_Track05_wav" + + +def create_mock_wav(output_fpath: str, duration_min: int) -> None: + duration_seconds = duration_min * 60 + sample_rate = 44100 + n_channels = 2 # Stereo + sampwidth = 2 # 2 bytes per sample (16-bit audio) + + # Generate a silent audio data array + num_samples = duration_seconds * sample_rate + audio_data = np.zeros((num_samples, n_channels), dtype=np.int16) + + # Create the WAV file + with wave.open(str(output_fpath), "w") as wav_file: + wav_file.setnchannels(n_channels) + wav_file.setsampwidth(sampwidth) + wav_file.setframerate(sample_rate) + wav_file.writeframes(audio_data.tobytes()) + + logging.info(f"Mock {duration_min}-minute WAV file '{output_fpath}' created successfully.") + + +def create_mock_midi(output_fpath: str) -> None: + # Create a new MIDI file with one track + mid = MidiFile() + track = MidiTrack() + mid.tracks.append(track) + + # Define a sequence of notes (time, type, note, velocity) + notes = [ + (0, "note_on", 60, 64), # C4 + (500, "note_off", 60, 64), + (0, "note_on", 62, 64), # D4 + (500, "note_off", 62, 64), + ] + + # Add the notes to the track + for time, type, note, velocity in notes: + track.append(Message(type, note=note, velocity=velocity, time=time)) + + # Save the MIDI file + mid.save(output_fpath) + + logging.info(f"Mock MIDI file '{output_fpath}' created successfully.") + + +def test_maestro_to_tf_example(tmp_path: pathlib.Path) -> None: + mock_maestro_home = tmp_path / "maestro" + mock_maestro_ext = mock_maestro_home / "2004" + mock_maestro_ext.mkdir(parents=True, exist_ok=True) + + create_mock_wav(str(mock_maestro_ext / (TRAIN_TRACK_ID.split("/")[1] + ".wav")), 3) + create_mock_midi(str(mock_maestro_ext / (TRAIN_TRACK_ID.split("/")[1] + ".midi"))) + + output_dir = tmp_path / "outputs" + output_dir.mkdir(parents=True, exist_ok=True) + + input_data: List[str] = [TRAIN_TRACK_ID] + with TestPipeline() as p: + ( + p + | "Create PCollection of track IDs" >> beam.Create([input_data]) + | "Create tf.Example" >> beam.ParDo(MaestroToTfExample(str(mock_maestro_home), download=False)) + | "Write to tfrecord" >> beam.ParDo(WriteBatchToTfRecord(str(output_dir))) + ) + + listdir = os.listdir(str(output_dir)) + assert len(listdir) == 1 + assert os.path.splitext(listdir[0])[-1] == ".tfrecord" + with open(os.path.join(str(output_dir), os.listdir(str(output_dir))[0]), "rb") as fp: + data = fp.read() + assert len(data) != 0 + + +def test_maestro_invalid_tracks(tmp_path: pathlib.Path) -> None: + mock_maestro_home = tmp_path / "maestro" + mock_maestro_ext = mock_maestro_home / "2004" + mock_maestro_ext.mkdir(parents=True, exist_ok=True) + + input_data = [(TRAIN_TRACK_ID, "train"), (VALID_TRACK_ID, "validation"), (TEST_TRACK_ID, "test")] + + for track_id, _ in input_data: + create_mock_wav(str(mock_maestro_ext / (track_id.split("/")[1] + ".wav")), 3) + + split_labels = set([e[1] for e in input_data]) + with TestPipeline() as p: + splits = ( + p + | "Create PCollection" >> beam.Create(input_data) + | "Tag it" >> beam.ParDo(MaestroInvalidTracks(str(mock_maestro_home))).with_outputs(*split_labels) + ) + + for split in split_labels: + ( + getattr(splits, split) + | f"Write {split} to text" + >> beam.io.WriteToText(str(tmp_path / f"output_{split}.txt"), shard_name_template="") + ) + + for track_id, split in input_data: + with open(tmp_path / f"output_{split}.txt", "r") as fp: + assert fp.read().strip() == track_id + + +def test_maestro_invalid_tracks_over_15_min(tmp_path: pathlib.Path) -> None: + """ + The track id used here is a real track id in maestro, and it is part of the train split, but we mock the data so as + not to store a large file in git, hence the variable name. + """ + + mock_maestro_home = tmp_path / "maestro" + mock_maestro_ext = mock_maestro_home / "2004" + mock_maestro_ext.mkdir(parents=True, exist_ok=True) + + mock_fpath = mock_maestro_ext / (GT_15M_TRACK_ID.split("/")[1] + ".wav") + create_mock_wav(str(mock_fpath), 16) + + input_data = [(GT_15M_TRACK_ID, "train")] + split_labels = set([e[1] for e in input_data]) + with TestPipeline() as p: + splits = ( + p + | "Create PCollection" >> beam.Create(input_data) + | "Tag it" >> beam.ParDo(MaestroInvalidTracks(str(mock_maestro_home))).with_outputs(*split_labels) + ) + + for split in split_labels: + ( + getattr(splits, split) + | f"Write {split} to text" + >> beam.io.WriteToText(str(tmp_path / f"output_{split}.txt"), shard_name_template="") + ) + + for _, split in input_data: + with open(tmp_path / f"output_{split}.txt", "r") as fp: + assert fp.read().strip() == "" + + +def test_maestro_create_input_data() -> None: + """ + A commuted metadata file is included in the repo for testing. mirdata references the metadata file to + populate the tracklist with metadata. Since the file is commuted to only the filenames referenced here, + we only consider these when testing the metadata. + """ + data = create_input_data(str(MAESTRO_TEST_DATA_PATH)) + assert len(data) + + test_fnames = {TRAIN_TRACK_ID, VALID_TRACK_ID, TEST_TRACK_ID, GT_15M_TRACK_ID} + splits = {d[1] for d in data if d[0].split(".")[0] in test_fnames} + assert splits == set(["train", "validation", "test"]) diff --git a/tests/resources/data/maestro/maestro-v2.0.0.json b/tests/resources/data/maestro/maestro-v2.0.0.json new file mode 100644 index 0000000..f81935b --- /dev/null +++ b/tests/resources/data/maestro/maestro-v2.0.0.json @@ -0,0 +1 @@ +[{"canonical_composer": "Fr\u00e9d\u00e9ric Chopin", "canonical_title": "12 Etudes, Op. 25", "split": "train", "year": 2004, "midi_filename": "2004/MIDI-Unprocessed_SMF_05_R1_2004_01_ORIG_MID--AUDIO_05_R1_2004_03_Track03_wav.midi", "audio_filename": "2004/MIDI-Unprocessed_SMF_05_R1_2004_01_ORIG_MID--AUDIO_05_R1_2004_03_Track03_wav.wav", "duration": 327.327025835}, {"canonical_composer": "Johann Sebastian Bach", "canonical_title": "French Suite No. 5 in G Major", "split": "train", "year": 2004, "midi_filename": "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_05_Track05_wav.midi", "audio_filename": "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_05_Track05_wav.wav", "duration": 968.256759019}, {"canonical_composer": "Sergei Rachmaninoff", "canonical_title": "Etude-Tableaux, Opus 33 No. 3 in C Minor", "split": "validation", "year": 2004, "midi_filename": "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_06_Track06_wav.midi", "audio_filename": "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_06_Track06_wav.wav", "duration": 267.438313619}, {"canonical_composer": "Sergei Rachmaninoff", "canonical_title": "Etude-Tableaux, Opus 39 No. 8 in D Minor", "split": "test", "year": 2004, "midi_filename": "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_08_Track08_wav.midi", "audio_filename": "2004/MIDI-Unprocessed_SMF_02_R1_2004_01-05_ORIG_MID--AUDIO_02_R1_2004_08_Track08_wav.wav", "duration": 190.852001088}] \ No newline at end of file diff --git a/tox.ini b/tox.ini index 68cdd27..e22a6af 100644 --- a/tox.ini +++ b/tox.ini @@ -35,7 +35,7 @@ commands = [testenv:manifest] deps = check-manifest skip_install = true -commands = check-manifest --ignore 'tests/*' +commands = check-manifest --ignore 'tests' [testenv:check-formatting] basepython = python3.10