<jupyter_start><jupyter_text>RealTrack data preparationThis notebook prepares the RealTrack (RT) training data. `INPUT_DIR` is expected to point to the `RealTracks` directory from Band-in-a-Box UltraPAK 2018, so that each subdirectory of `INPUT_DIR` contains the WMA files for a single RealTrack.The code creates a `wav_16kHz` directory containing the 8-second training segments, a `metadata.json` file with information about each segment, and `pairs_train`, `pairs_val` and `pairs_test` files listing pairs of audio file paths. Note that `pairs_test` is not used in the paper.Copyright 2020 InterDigital R&D and Télécom Paris.  Author: <NAME><jupyter_code>import argparse
import collections
import concurrent.futures as cf
import glob
import hashlib
import json
import os
import random
import regex as re
import sys

import IPython.display as ipd
import itertools
import librosa
import matplotlib.pyplot as plt
from natsort import natsorted, ns
import numpy as np
import pandas as pd
import pydub
import pysndfx
from pydub import silence
from sklearn.model_selection import train_test_split
import soundfile as sf
from tqdm.auto import tqdm
INPUT_DIR = 'RealTracks'
OUTPUT_DIR = 'wav_16kHz'

# Load data augmentation parameters from metadata_ref.json instead of sampling them randomly.
# Set to True to reproduce the dataset from the paper. Set to False if you want to use your own data.
USE_REF_METADATA = True

MAX_FINAL_CHUNK_LEN = 8000  # 8 s
MAX_TRANSPOSE = 4  # maximum number of transposition steps
MAX_RESAMPLE = 2 ** (4 / 12)
MAX_CHUNK_LEN = int(MAX_FINAL_CHUNK_LEN * MAX_RESAMPLE / 1000 + 0.5) * 1000  # 10 s
MAX_CHUNKS = 20 * 60 * 1000 // MAX_CHUNK_LEN  # 20 min
SR = 16000
def random_fx(rng):
    chain = pysndfx.AudioEffectsChain()
    for _ in range(rng.choice([0, 1, 2, 3], p=[0.35, 0.4, 0.2, 0.05])):
        effect = rng.choice([
            lambda c: c.overdrive(gain=rng.uniform(10, 40)),
            lambda c: c.phaser(gain_in=rng.uniform(0.6, 0.9),
                               gain_out=rng.uniform(0.66, 0.85),
                               delay=rng.power(0.4) * 3 + 1,
                               decay=rng.uniform(0.2, 0.45),
                               speed=rng.uniform(0.5, 2),
                               triangular=rng.choice([True, False])),
            lambda c: c.gain(-3).reverb(),
            lambda c: c.tremolo(freq=rng.power(0.5) * 14 + 1,
                                depth=rng.uniform(20, 60))
        ])
        effect(chain)
    return chain
def pydub_to_numpy(audio, sr=SR):
    samples = audio.set_channels(1).set_frame_rate(sr).get_array_of_samples()
    array = np.array(samples)
    src_dtype = array.dtype
    array = array.astype(np.float32)
    if len(samples) > 0:
        array /= np.iinfo(src_dtype).max
    return array
def process_dir(args):
    dir_path = args
    dir_name = os.path.basename(dir_path)
    out_dir = os.path.join(OUTPUT_DIR, dir_name)
    os.makedirs(out_dir)
    
    if USE_REF_METADATA:
        with open('metadata_ref.json') as f:
            metadata_ref_raw = json.load(f)
        metadata_ref = collections.OrderedDict()
        for item in metadata_ref_raw:
            if item['style'] == dir_name:
                ref_key = tuple(os.path.splitext(os.path.basename(segment['path']))[0]
                                for segment in item['segments'])
                assert ref_key not in metadata_ref
                metadata_ref[ref_key] = item
    else:
        seed = int.from_bytes(hashlib.sha512(dir_name.encode()).digest(), 'big')
        rng = np.random.default_rng(seed=seed)

    chunks = []
    for path in glob.glob(os.path.join(dir_path, '*.wma')):
        filename = os.path.basename(path)
        name, _ = os.path.splitext(filename)
        audio = pydub.AudioSegment.from_file(path)
        for t in range(0, len(audio), MAX_CHUNK_LEN):
            end = t + MAX_CHUNK_LEN
            chunks.append(argparse.Namespace(
                audio=audio, start=t, end=end,
                name=f'{name}.{t // 1000}-{end // 1000}',
                src_path=os.path.join(dir_name, filename)))
    chunks.sort(key=lambda chk: (chk.src_path, chk.name))

    if USE_REF_METADATA:
        chunks_by_name = {chk.name: chk for chk in chunks}
        chunks = [chunks_by_name[name] for key in metadata_ref for name in key]
    else:
        rng.shuffle(chunks)
        chunks = chunks[:MAX_CHUNKS]

    metadata = []
    for chunk_pair in zip(chunks[::2], chunks[1::2]):
        meta = {
            'style': dir_name
        }

        if USE_REF_METADATA:
            ref_key = tuple(chk.name for chk in chunk_pair)
            meta['resample_rate'] = metadata_ref[ref_key]['resample_rate']
            effects = pysndfx.AudioEffectsChain()
            effects.command = metadata_ref[ref_key]['effects']
        else:
            # Pick a resampling rate so as to transpose by a whole number of steps
            meta['resample_rate'] = 1 / (2 ** (rng.choice(np.arange(-MAX_TRANSPOSE, MAX_TRANSPOSE + 1)) / 12))

            # Pick random effects
            effects = random_fx(rng)
        meta['effects'] = effects.command

        meta['segments'] = []
        for chunk in chunk_pair:
            audio = pydub_to_numpy(chunk.audio[chunk.start:chunk.end])
            
            # Apply effects
            audio = librosa.resample(audio, SR, SR * meta['resample_rate'])
            audio = effects(audio, sample_in=SR)

            # Clip to maximum length
            audio = audio[:MAX_FINAL_CHUNK_LEN * SR // 1000]

            out_path = os.path.join(out_dir, f'{chunk.name}.wav')
            sf.write(out_path, audio, SR, subtype='PCM_24')
            
            meta['segments'].append({
                'src_path': chunk.src_path,
                'path': os.path.relpath(out_path, OUTPUT_DIR),
                'index': chunk.start // MAX_CHUNK_LEN
            })
        metadata.append(meta)

    return dir_name, metadata

dir_paths = sorted(set(os.path.dirname(p) for p in glob.glob(os.path.join(INPUT_DIR, '*', '*.wma'))))
metadata = []
with cf.ProcessPoolExecutor(16) as executor:
    for style, metadata_chunk in tqdm(executor.map(process_dir, dir_paths), total=len(dir_paths)):
        metadata.extend(metadata_chunk)
        print(style, 'OK')
metadata = natsorted(metadata, key=lambda m: tuple(s['path'] for s in m['segments']), alg=ns.PATH)
with open('metadata.json', 'w') as f:
    json.dump(metadata, f)
out_path_pairs = [(os.path.join(OUTPUT_DIR, a['path']), os.path.join(OUTPUT_DIR, b['path']))
                  for m in metadata for a, b in [m['segments']]]
np.random.seed(42)
np.random.shuffle(out_path_pairs)
np.random.seed(42)
train_pairs, test_pairs = train_test_split(out_path_pairs, test_size=0.01)
train_pairs, val_pairs = train_test_split(train_pairs, test_size=500)
np.random.shuffle(train_pairs)

def write_tuples(tuples, path, shuffle_items=False):
    with open(path, 'w') as f:
        for tup in tuples:
            if shuffle_items:
                tup = np.random.choice(tup, size=len(tup), replace=False)
            print(*tup, sep='\t', file=f)

write_tuples(train_pairs, 'pairs_train', shuffle_items=True)
write_tuples(val_pairs, 'pairs_val')
write_tuples(test_pairs, 'pairs_test')
!wc -l pairs_*
len(dir_paths)  # 1526<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11310420375188279, 'fraction_numerical': 0.020128714227029987, 'mean_word_length': 2.688888888888889, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '566650', 'license': 'Apache-2.0', 'n_tokens_mistral': 2542, 'n_tokens_neox': 2422, 'n_words': 587, 'path': 'data/rt/prepare.ipynb', 'repo_name': '18praveenb/ss-vq-vae', 'size': '7309'}