added model weights

.gitattributes

@@ -1 +1,2 @@
 *.wav filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,5 +1,8 @@
 __pycache__
 .DS_Store
-data
+data/samples
+data/samples-backup.zip
+data/samples-backup.zip
+data/songs.csv
 logs
 gradio_cached_examples

app.py

@@ -4,15 +4,17 @@ import numpy as np
 import torch
 from preprocessing.preprocess import AudioPipeline
 from preprocessing.preprocess import AudioPipeline
-from dancer_net.dancer_net import ShortChunkCNN
+from models.residual import ResidualDancer
 import os
 import json
 from functools import cache
 import pandas as pd
 
+
+
 @cache
-def get_model(device) -> tuple[ShortChunkCNN, np.ndarray]:
-    model_path = "logs/20221226-230930"
+def get_model(device) -> tuple[ResidualDancer, np.ndarray]:
+    model_path = "models/weights/ResidualDancer"
     weights = os.path.join(model_path, "dancer_net.pt")
     config_path = os.path.join(model_path, "config.json")
 
@@ -20,7 +22,7 @@ def get_model(device) -> tuple[ShortChunkCNN, np.ndarray]:
         config = json.load(f)
     labels = np.array(sorted(config["classes"]))
 
-    model = ShortChunkCNN(n_class=len(labels))
+    model = ResidualDancer(n_classes=len(labels))
     model.load_state_dict(torch.load(weights))
     model = model.to(device).eval()
     return model, labels

data/dance_mapping.csv

@@ -0,0 +1,48 @@
+id,name
+SWZ,Slow Waltz
+CSW,Cross-step Waltz
+CFT,Castle Foxtrot
+SFT,Slow Foxtrot
+TGO,Tango (Ballroom)
+PBD,Peabody
+VWZ,Viennese Waltz
+QST,Quickstep
+BOL,Bolero
+CHA,Cha Cha
+MBO,Mambo
+JIV,Jive
+RMB,Rumba
+ECS,East Coast Swing
+WCS,West Coast Swing
+HST,Hustle
+MRG,Merengue
+PDL,Paso Doble
+SMB,Samba
+PLK,Polka
+SLS,Salsa
+BCH,Bachata
+NC2,Night Club Two Step
+C2S,Country Two Step
+CMB,Cumbia
+LHP,Lindy Hop
+CST,Charleston
+CSG,Carolina Shag
+CLS,Collegiate Shag
+ATN,Argentine Tango
+TGV,Tango Vals
+NTN,Neo Tango
+MGA,Milonga
+BSN,Bossa Nova
+JSW,Jump Swing
+BLU,Blues
+MWT,Motown
+BBA,Balboa
+JAZ,Jazz
+CNT,Contemporary
+BLT,Ballet
+BDW,Broadway
+TAP,Tap
+HHP,Hip-Hop
+BWD,Bollywood
+DSC,Disco
+FST,Freestyle

main.py

@@ -1,46 +0,0 @@
-import torchaudio
-from preprocessing.preprocess import AudioPipeline
-from dancer_net.dancer_net import ShortChunkCNN
-import torch
-import numpy as np
-import os
-import json
-
-if __name__ == "__main__":
-
-    audio_file = "data/samples/mzm.iqskzxzx.aac.p.m4a.wav"
-    seconds = 6
-    model_path = "logs/20221226-230930"
-    weights = os.path.join(model_path, "dancer_net.pt")
-    config_path = os.path.join(model_path, "config.json")
-    device = "mps"
-    threshold = 0.5
-
-    with open(config_path) as f:
-        config = json.load(f)
-    labels = np.array(sorted(config["classes"]))
-
-    audio_pipeline = AudioPipeline()
-    waveform, sample_rate = torchaudio.load(audio_file)
-    waveform = waveform[:, :seconds * sample_rate]
-    spectrogram = audio_pipeline(waveform)
-    spectrogram = spectrogram.unsqueeze(0).to(device)
-
-    model = ShortChunkCNN(n_class=len(labels))
-    model.load_state_dict(torch.load(weights))
-    model = model.to(device).eval()
-
-    with torch.no_grad():
-        results = model(spectrogram)
-    results = results.squeeze(0).detach().cpu().numpy()
-    results = results > threshold
-    results = labels[results]
-    print(results)
-
-
-
-
-    
-
-
-

dancer_net/dancer_net.py → models/residual.py

@@ -1,16 +1,12 @@
-import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torchaudio import transforms as taT, functional as taF
 
+# Architecture based on: https://github.com/minzwon/sota-music-tagging-models/blob/36aa13b7205ff156cf4dcab60fd69957da453151/training/model.py
 
-
-DEVICE = "mps"
-class ShortChunkCNN(nn.Module):
+class ResidualDancer(nn.Module):
     def __init__(self,
                 n_channels=128,
-                sample_rate=16000,
-                n_class=50):
+                n_classes=50):
         super().__init__()
 
         # Spectrogram
@@ -18,19 +14,19 @@ class ShortChunkCNN(nn.Module):
 
         # CNN
         self.res_layers = nn.Sequential(
-            Res_2d(1, n_channels, stride=2),
-            Res_2d(n_channels, n_channels, stride=2),
-            Res_2d(n_channels, n_channels*2, stride=2),
-            Res_2d(n_channels*2, n_channels*2, stride=2),
-            Res_2d(n_channels*2, n_channels*2, stride=2),
-            Res_2d(n_channels*2, n_channels*2, stride=2),
-            Res_2d(n_channels*2, n_channels*4, stride=2)
+            ResBlock(1, n_channels, stride=2),
+            ResBlock(n_channels, n_channels, stride=2),
+            ResBlock(n_channels, n_channels*2, stride=2),
+            ResBlock(n_channels*2, n_channels*2, stride=2),
+            ResBlock(n_channels*2, n_channels*2, stride=2),
+            ResBlock(n_channels*2, n_channels*2, stride=2),
+            ResBlock(n_channels*2, n_channels*4, stride=2)
         )
 
         # Dense
         self.dense1 = nn.Linear(n_channels*4, n_channels*4)
         self.bn = nn.BatchNorm1d(n_channels*4)
-        self.dense2 = nn.Linear(n_channels*4, n_class)
+        self.dense2 = nn.Linear(n_channels*4, n_classes)
         self.dropout = nn.Dropout(0.3)
 
     def forward(self, x):
@@ -56,7 +52,7 @@ class ShortChunkCNN(nn.Module):
         return x
 
 
-class Res_2d(nn.Module):
+class ResBlock(nn.Module):
     def __init__(self, input_channels, output_channels, shape=3, stride=2):
         super().__init__()
         # convolution

models/weights/ResidualDancer/config.json

@@ -0,0 +1,24 @@
+{
+    "classes": [
+        "ATN",
+        "BBA",
+        "BCH",
+        "BLU",
+        "CHA",
+        "CMB",
+        "CSG",
+        "ECS",
+        "HST",
+        "JIV",
+        "LHP",
+        "QST",
+        "RMB",
+        "SFT",
+        "SLS",
+        "SMB",
+        "SWZ",
+        "TGO",
+        "VWZ",
+        "WCS"
+    ]
+}

models/weights/ResidualDancer/dancer_net.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1888558eed82a5d99ac1dab55969a9ea36455d11a9370355d1f2b984598d30ff
+size 48453416

train.py

@@ -13,10 +13,30 @@ from sklearn.model_selection import KFold
 from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score
 from preprocessing.dataset import SongDataset
 from preprocessing.preprocess import get_examples
-from dancer_net.dancer_net import ShortChunkCNN
+from models.residual import ResidualDancer
 
 DEVICE = "mps"
 SEED = 42
+TARGET_CLASSES = ['ATN',
+        'BBA',
+        'BCH',
+        'BLU',
+        'CHA',
+        'CMB',
+        'CSG',
+        'ECS',
+        'HST',
+        'JIV',
+        'LHP',
+        'QST',
+        'RMB',
+        'SFT',
+        'SLS',
+        'SMB',
+        'SWZ',
+        'TGO',
+        'VWZ',
+        'WCS']
 
 def get_timestamp() -> str:
     return datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S")
@@ -115,28 +135,8 @@ def train(
         
 
 def cross_validation(seed=42, batch_size=64, k=5, device="mps"):
-    target_classes = ['ATN',
-    'BBA',
-    'BCH',
-    'BLU',
-    'CHA',
-    'CMB',
-    'CSG',
-    'ECS',
-    'HST',
-    'JIV',
-    'LHP',
-    'QST',
-    'RMB',
-    'SFT',
-    'SLS',
-    'SMB',
-    'SWZ',
-    'TGO',
-    'VWZ',
-    'WCS']
     df = pd.read_csv("data/songs.csv")
-    x,y = get_examples(df, "data/samples",class_list=target_classes)
+    x,y = get_examples(df, "data/samples",class_list=TARGET_CLASSES)
     
     dataset = SongDataset(x,y)
     splits=KFold(n_splits=k,shuffle=True,random_state=seed)
@@ -149,7 +149,7 @@ def cross_validation(seed=42, batch_size=64, k=5, device="mps"):
         train_loader = DataLoader(dataset, batch_size=batch_size, sampler=train_sampler)
         test_loader = DataLoader(dataset, batch_size=batch_size, sampler=test_sampler)
         n_classes = len(y[0])
-        model = ShortChunkCNN(n_class=n_classes).to(device)
+        model = ResidualDancer(n_classes=n_classes).to(device)
         model, _ = train(model,train_loader, epochs=2, device=device)
         val_metrics = evaluate(model, test_loader, nn.BCELoss())
         metrics.append(val_metrics)
@@ -164,28 +164,9 @@ def cross_validation(seed=42, batch_size=64, k=5, device="mps"):
 
 
 def train_model():
-    target_classes = ['ATN',
-        'BBA',
-        'BCH',
-        'BLU',
-        'CHA',
-        'CMB',
-        'CSG',
-        'ECS',
-        'HST',
-        'JIV',
-        'LHP',
-        'QST',
-        'RMB',
-        'SFT',
-        'SLS',
-        'SMB',
-        'SWZ',
-        'TGO',
-        'VWZ',
-        'WCS']
+
     df = pd.read_csv("data/songs.csv")
-    x,y = get_examples(df, "data/samples",class_list=target_classes)
+    x,y = get_examples(df, "data/samples",class_list=TARGET_CLASSES)
     dataset = SongDataset(x,y)
     train_count = int(len(dataset) * 0.9)
     datasets = random_split(dataset, [train_count, len(dataset) - train_count], torch.Generator().manual_seed(SEED))
@@ -193,7 +174,7 @@ def train_model():
     train_data, val_data = data_loaders
     example_spec, example_label = dataset[0]
     n_classes = len(example_label)
-    model = ShortChunkCNN(n_class=n_classes).to(DEVICE)
+    model = ResidualDancer(n_classes=n_classes).to(DEVICE)
     model, metrics = train(model,train_data, val_data, epochs=3, device=DEVICE)
     
     log_dir = os.path.join(
@@ -201,11 +182,11 @@ def train_model():
     )
     os.makedirs(log_dir, exist_ok=True)
     
-    torch.save(model.state_dict(), os.path.join(log_dir, "dancer_net.pt"))
+    torch.save(model.state_dict(), os.path.join(log_dir, "residual_dancer.pt"))
     metrics = pd.DataFrame(metrics)
     metrics.to_csv(os.path.join(log_dir, "metrics.csv"))
     config = {
-        "classes": target_classes
+        "classes": TARGET_CLASSES
     }
     with open(os.path.join(log_dir, "config.json")) as f:
         json.dump(config, f)