added slakh processing, tests, and test data, tox passing.

.github/workflows/tox.yml

@@ -30,13 +30,13 @@ jobs:
           python-version: ${{ matrix.py }}
       - uses: actions/checkout@v3
       - name: Install soundlibs Ubuntu
-        run: sudo apt-get update && sudo apt-get install --no-install-recommends -y --fix-missing pkg-config libsndfile1 sox
+        run: sudo apt-get update && sudo apt-get install --no-install-recommends -y --fix-missing pkg-config libsndfile1 sox ffmpeg
         if: matrix.os == 'ubuntu-latest'
       - name: Install soundlibs MacOs
-        run: brew install libsndfile llvm libomp sox
+        run: brew install libsndfile llvm libomp sox ffmpeg
         if: matrix.os == 'macos-latest-xlarge'
       - name: Install soundlibs Windows
-        run: choco install libsndfile sox.portable
+        run: choco install libsndfile sox.portable flac ffmpeg
         if: matrix.os == 'windows-latest'
       - name: Upgrade pip
         run: python -m pip install -U pip

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 *.mid *.midi *.flac *.yaml *.json
 recursive-include basic_pitch *.py *.md
 recursive-include basic_pitch/saved_models *.index *.pb variables.data* *.mlmodel *.json *.onnx *.tflite *.bin

basic_pitch/data/datasets/guitarset.py

@@ -136,17 +136,19 @@ def create_input_data(
     if seed:
         random.seed(seed)
 
-    def determine_split() -> str:
-        partition = random.uniform(0, 1)
-        if partition < validation_bound:
+    def determine_split(index: int) -> str:
+        if index < len(track_ids) * validation_bound:
             return "train"
-        if partition < test_bound:
+        elif index < len(track_ids) * test_bound:
             return "validation"
-        return "test"
+        else:
+            return "test"
 
     guitarset = mirdata.initialize("guitarset")
+    track_ids = guitarset.track_ids
+    random.shuffle(track_ids)
 
-    return [(track_id, determine_split()) for track_id in guitarset.track_ids]
+    return [(track_id, determine_split(i)) for i, track_id in enumerate(track_ids)]
 
 
 def main(known_args: argparse.Namespace, pipeline_args: List[str]) -> None:

basic_pitch/data/datasets/ikala.py

@@ -138,21 +138,17 @@ def process(self, element: List[str], *args: Tuple[Any, Any], **kwargs: Dict[str
 def create_input_data(train_percent: float, seed: Optional[int] = None) -> List[Tuple[str, str]]:
     assert train_percent < 1.0, "Don't over allocate the data!"
 
-    # Test percent is 1 - train - validation
-    validation_bound = train_percent
-
     if seed:
         random.seed(seed)
 
-    def determine_split() -> str:
-        partition = random.uniform(0, 1)
-        if partition < validation_bound:
-            return "train"
-        return "validation"
-
     ikala = mirdata.initialize("ikala")
+    track_ids = ikala.track_ids
+    random.shuffle(track_ids)
+
+    def determine_split(index: int) -> str:
+        return "train" if index < len(track_ids) * train_percent else "validation"
 
-    return [(track_id, determine_split()) for track_id in ikala.track_ids]
+    return [(track_id, determine_split(i)) for i, track_id in enumerate(track_ids)]
 
 
 def main(known_args: argparse.Namespace, pipeline_args: List[str]) -> None:

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)