update splits, reqs

scripts/utils/split.py

@@ -1,8 +1,12 @@
 from pathlib import Path
 import random
 import shutil
+import os
+import json 
 
 import argbind
+from tqdm import tqdm
+from tqdm.contrib.concurrent import thread_map
 
 from audiotools.core import util
 
@@ -12,8 +16,13 @@ def train_test_split(
     audio_folder: str = ".", 
     test_size: float = 0.2,
     seed: int = 42,
+    pattern: str = "**/*.mp3",
 ):
-    audio_files = util.find_audio(audio_folder)
+    print(f"finding audio")
+
+    audio_folder = Path(audio_folder)
+    audio_files = list(tqdm(audio_folder.glob(pattern)))
+    print(f"found {len(audio_files)} audio files")
     
     # split according to test_size
     n_test = int(len(audio_files) * test_size)
@@ -37,10 +46,14 @@ def train_test_split(
     for split, files in (
         ("train", train_files), ("test", test_files)
     ):
-        for file in files:
-            out_file = Path(file).parent / split / Path(file).name
+        for file in tqdm(files):
+            out_file = audio_folder.parent / f"{audio_folder.name}-{split}" / Path(file).name
             out_file.parent.mkdir(exist_ok=True, parents=True)
-            shutil.copy(file, out_file)
+            os.symlink(file, out_file)
+
+        # save split as json
+        with open(Path(audio_folder) / f"{split}.json", "w") as f:
+            json.dump([str(f) for f in files], f)
     
 
     

setup.py

@@ -39,6 +39,7 @@ setup(
         "google-cloud-logging==2.2.0",
         "einops",
         # "frechet_audio_distance", 
-        "gradio"
+        "gradio", 
+        "tensorboardX",
     ],
 )