updated production weights

app.py

@@ -2,18 +2,21 @@ from pathlib import Path
 import gradio as gr
 import numpy as np
 import os
+import pandas as pd
 from functools import cache
 from pathlib import Path
-from models.audio_spectrogram_transformer import AST, ASTExtractorWrapper
+from models.residual import ResidualDancer
 from models.training_environment import TrainingEnvironment
+from preprocessing.pipelines import SpectrogramProductionPipeline
 import torch
 from torch import nn
 import yaml
 import torchaudio
 
-CONFIG_FILE = Path("models/config/train_local.yaml")
-MODEL_CLS = AST
-EXTRACTOR = ASTExtractorWrapper
+CONFIG_FILE = Path("models/config/production.yaml")
+MODEL_CLS = ResidualDancer
+
+DANCE_MAPPING_FILE = Path("data/dance_mapping.csv")
 
 
 class DancePredictor:
@@ -22,7 +25,7 @@ class DancePredictor:
         weight_path: str,
         labels: list[str],
         expected_duration=6,
-        threshold=0.5,
+        threshold=0.1,
         resample_frequency=16000,
         device="cpu",
     ):
@@ -35,11 +38,13 @@ class DancePredictor:
         self.labels = np.array(labels)
         self.device = device
         self.model = self.get_model(weight_path)
-        self.extractor = ASTExtractorWrapper()
+        self.extractor = SpectrogramProductionPipeline()
 
     def get_model(self, weight_path: str) -> nn.Module:
         weights = torch.load(weight_path, map_location=self.device)["state_dict"]
-        model = AST(self.labels).to(self.device)
+        n_classes = len(self.labels)
+        # NOTE: Channels are not taken into account
+        model = ResidualDancer(n_classes=n_classes).to(self.device)
         for key in list(weights):
             weights[
                 key.replace(
@@ -56,10 +61,12 @@ class DancePredictor:
             config = yaml.safe_load(f)
         weight_path = config["checkpoint"]
         labels = sorted(config["dance_ids"])
-        expected_duration = 6
-        threshold = 0.5
-        resample_frequency = 16000
-        device = "mps"
+        dance_mapping = get_dance_mapping(DANCE_MAPPING_FILE)
+        labels = [dance_mapping[label] for label in labels]
+        expected_duration = config.get("expected_duration", 6)
+        threshold = config.get("threshold", 0.1)
+        resample_frequency = config.get("resample_frequency", 16000)
+        device = config.get("device", "cpu")
         return DancePredictor(
             weight_path,
             labels,
@@ -81,9 +88,6 @@ class DancePredictor:
         waveform = torchaudio.functional.resample(
             waveform, sample_rate, self.resample_frequency
         )
-        waveform = waveform[
-            :, : self.resample_frequency * self.expected_duration
-        ]  # TODO PAD
         features = self.extractor(waveform)
         features = features.unsqueeze(0).to(self.device)
         results = self.model(features)
@@ -103,7 +107,15 @@ def get_model(config_path: str) -> DancePredictor:
     return model
 
 
+@cache
+def get_dance_mapping(mapping_file: str) -> dict[str, str]:
+    mapping_df = pd.read_csv(mapping_file)
+    return {row["id"]: row["name"] for _, row in mapping_df.iterrows()}
+
+
 def predict(audio: tuple[int, np.ndarray]) -> list[str]:
+    if audio is None:
+        return "Dance Not Found"
     sample_rate, waveform = audio
 
     model = get_model(CONFIG_FILE)
@@ -116,7 +128,7 @@ def demo():
     description = "What should I dance to this song? Pass some audio to the Dance Classifier find out!"
     song_samples = Path(os.path.dirname(__file__), "assets", "song-samples")
     example_audio = [
-        str(song) for song in song_samples.iterdir() if song.name[0] != "."
+        str(song) for song in song_samples.iterdir() if not song.name.startswith(".")
     ]
     all_dances = get_model(CONFIG_FILE).labels
 

assets/song-samples/besame_mucho.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14ccffab50d9119ec5250fc84e09542dbbf350450102c108ab61846a3c3031c8
+size 5290062

models/config/production.yaml

@@ -0,0 +1,20 @@
+checkpoint: models/weights/ResidualDancer/weights.ckpt
+device: cpu
+seed: 42
+dance_ids: &dance_ids
+  - BCH
+  - CHA
+  - JIV
+  - ECS
+  - QST
+  - RMB
+  - SFT
+  - SLS
+  - SMB
+  - SWZ
+  - TGO
+  - VWZ
+  - WCS
+
+model:
+  n_channels: 128

models/config/train_local.yaml

@@ -1,5 +1,5 @@
-training_fn: wav2vec2.train_huggingface
-checkpoint: lightning_logs/version_172/checkpoints/epoch=3-step=4572.ckpt
+training_fn: residual.train_residual_dancer
+checkpoint: lightning_logs/version_176/checkpoints/epoch=12-step=40404.ckpt
 device: mps
 seed: 42
 dance_ids: &dance_ids
@@ -24,10 +24,10 @@ data_module:
   test_proportion: 0.2
 
 datasets:
-  # preprocessing.dataset.BestBallroomDataset:
-  #   audio_dir: data/ballroom-songs
-  #   class_list: *dance_ids
-  #   audio_window_jitter: 0.7
+  preprocessing.dataset.BestBallroomDataset:
+    audio_dir: data/ballroom-songs
+    class_list: *dance_ids
+    audio_window_jitter: 0.7
 
   preprocessing.dataset.Music4DanceDataset:
     song_data_path: data/songs_cleaned.csv

models/residual.py

@@ -110,7 +110,7 @@ def train_residual_dancer(config: dict):
     TARGET_CLASSES = config["dance_ids"]
     DEVICE = config["device"]
     SEED = config["seed"]
-    torch.set_float32_matmul_precision('medium')
+    torch.set_float32_matmul_precision("medium")
     pl.seed_everything(SEED, workers=True)
     feature_extractor = SpectrogramTrainingPipeline(**config["feature_extractor"])
     dataset = get_datasets(config["datasets"], feature_extractor)
@@ -123,7 +123,7 @@ def train_residual_dancer(config: dict):
     train_env = TrainingEnvironment(model, criterion, config)
     callbacks = [
         # cb.LearningRateFinder(update_attr=True),
-        cb.EarlyStopping("val/loss", patience=5),
+        cb.EarlyStopping("val/loss", patience=1),
         cb.StochasticWeightAveraging(1e-2),
         cb.RichProgressBar(),
         cb.DeviceStatsMonitor(),

models/weights/ResidualDancer/weights.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e107090ff62ac0b79f4f40271e8b1dd6c3d10d8146264ec49df3c8febe99aa23
-size 193651217
+oid sha256:90a58841ce4f40f2981227b63dd848e474e8868795a57da84053e3281c4889c7
+size 193643085

preprocessing/dataset.py

@@ -78,8 +78,8 @@ class SongDataset(Dataset):
             return waveform, dance_labels
         else:
             # WARNING: Could cause train/test split leak
-            return self[idx-1]
-
+            print("Invalid output, trying next index...")
+            return self[idx - 1]
 
     def _idx2audio_idx(self, idx: int) -> int:
         return self._get_audio_loc_from_idx(idx)[0]
@@ -424,3 +424,7 @@ def record_audio_durations(folder: str):
 
     with open(os.path.join(folder, "audio_durations.json"), "w") as f:
         json.dump(durations, f)
+
+
+class GTZAN:
+    pass

preprocessing/pipelines.py

@@ -74,6 +74,21 @@ class SpectrogramTrainingPipeline(WaveformTrainingPipeline):
         return spec
 
 
+class SpectrogramProductionPipeline(torch.nn.Module):
+    def __init__(self, sample_rate=16000, expected_duration=6, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.preprocess_waveform = WaveformPreprocessing(
+            sample_rate * expected_duration
+        )
+        self.audio_to_spectrogram = AudioToSpectrogram(
+            sample_rate=sample_rate,
+        )
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        waveform = self.preprocess_waveform(waveform)
+        return self.audio_to_spectrogram(waveform)
+
+
 class WaveformPreprocessing(torch.nn.Module):
     def __init__(self, expected_sample_length: int):
         super().__init__()