basic_pitch/data/datasets/maestro.py

#!/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)