def main()

in tools/automix.py [0:0]


def main():
    random.seed(4321)
    if OUTPATH.exists():
        shutil.rmtree(OUTPATH)
    OUTPATH.mkdir(exist_ok=True, parents=True)
    (OUTPATH / 'train').mkdir(exist_ok=True, parents=True)
    (OUTPATH / 'valid').mkdir(exist_ok=True, parents=True)
    out = OUTPATH / 'train'

    dset = get_musdb_dataset()
    # dset2 = get_wav_dataset()
    # dset3 = get_musdb_dataset('test')
    dset2 = None
    dset3 = None
    pendings = []
    copies = 6
    copies_rej = 2

    with ProcessPoolExecutor(20) as pool:
        for index in range(len(dset)):
            pendings.append(pool.submit(analyse_track, dset, index))

        if dset2:
            for index in range(len(dset2)):
                pendings.append(pool.submit(analyse_track, dset2, index))
        if dset3:
            for index in range(len(dset3)):
                pendings.append(pool.submit(analyse_track, dset3, index))

        catalog = []
        rej = 0
        for pending in tqdm.tqdm(pendings, ncols=120):
            spec, track = pending.result()
            if spec is not None:
                catalog.append(spec)
            else:
                mix = track.sum(0)
                for copy in range(copies_rej):
                    folder = out / f'rej_{rej}_{copy}'
                    folder.mkdir()
                    save_audio(mix, folder / "mixture.wav", SR)
                    for stem, source in zip(track, SOURCES):
                        save_audio(stem, folder / f"{source}.wav", SR, clip='clamp')
                    rej += 1

    for copy in range(copies):
        for index in range(len(catalog)):
            track, origs = build_track(index, catalog)
            mix = track.sum(0)
            mx = mix.abs().max()
            scale = max(1, 1.01 * mx)
            mix = mix / scale
            track = track / scale
            folder = out / f'{copy}_{index}'
            folder.mkdir()
            save_audio(mix, folder / "mixture.wav", SR)
            for stem, source, orig in zip(track, SOURCES, origs):
                save_audio(stem, folder / f"{source}.wav", SR, clip='clamp')