Spaces:
Runtime error
Runtime error
File size: 4,695 Bytes
c914273 4b8361a 42c4703 c914273 42c4703 4b8361a c914273 42c4703 c914273 42c4703 4b8361a 42c4703 4b8361a c914273 4b8361a 42c4703 4b8361a 42c4703 4b8361a 42c4703 4b8361a 42c4703 4b8361a 42c4703 c914273 42c4703 4b8361a c914273 42c4703 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
from pathlib import Path
import gradio as gr
import numpy as np
import os
from functools import cache
from pathlib import Path
from models.audio_spectrogram_transformer import AST, ASTExtractorWrapper
from models.training_environment import TrainingEnvironment
import torch
from torch import nn
import yaml
import torchaudio
CONFIG_FILE = Path("models/config/train_local.yaml")
MODEL_CLS = AST
EXTRACTOR = ASTExtractorWrapper
class DancePredictor:
def __init__(
self,
weight_path: str,
labels: list[str],
expected_duration=6,
threshold=0.5,
resample_frequency=16000,
device="cpu",
):
super().__init__()
self.expected_duration = expected_duration
self.threshold = threshold
self.resample_frequency = resample_frequency
self.labels = np.array(labels)
self.device = device
self.model = self.get_model(weight_path)
self.extractor = ASTExtractorWrapper()
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)
for key in list(weights):
weights[
key.replace(
"model.",
"",
)
] = weights.pop(key)
model.load_state_dict(weights, strict=False)
return model.to(self.device).eval()
@classmethod
def from_config(cls, config_path: str) -> "DancePredictor":
with open(config_path, "r") as f:
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"
return DancePredictor(
weight_path,
labels,
expected_duration,
threshold,
resample_frequency,
device,
)
@torch.no_grad()
def __call__(self, waveform: np.ndarray, sample_rate: int) -> dict[str, float]:
if waveform.ndim == 1:
waveform = np.stack([waveform, waveform]).T
waveform = torch.from_numpy(waveform.T)
waveform = torchaudio.functional.apply_codec(
waveform, sample_rate, "wav", channels_first=True
)
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)
results = nn.functional.softmax(results.squeeze(0), dim=0)
results = results.detach().cpu().numpy()
result_mask = results > self.threshold
probs = results[result_mask]
dances = self.labels[result_mask]
return {dance: float(prob) for dance, prob in zip(dances, probs)}
@cache
def get_model(config_path: str) -> DancePredictor:
model = DancePredictor.from_config(config_path)
return model
def predict(audio: tuple[int, np.ndarray]) -> list[str]:
sample_rate, waveform = audio
model = get_model(CONFIG_FILE)
results = model(waveform, sample_rate)
return results if len(results) else "Dance Not Found"
def demo():
title = "Dance Classifier"
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] != "."
]
all_dances = get_model(CONFIG_FILE).labels
recording_interface = gr.Interface(
fn=predict,
description="Record at least **6 seconds** of the song.",
inputs=gr.Audio(source="microphone", label="Song Recording"),
outputs=gr.Label(label="Dances"),
examples=example_audio,
)
uploading_interface = gr.Interface(
fn=predict,
inputs=gr.Audio(label="Song Audio File"),
outputs=gr.Label(label="Dances"),
examples=example_audio,
)
with gr.Blocks() as app:
gr.Markdown(f"# {title}")
gr.Markdown(description)
gr.TabbedInterface(
[uploading_interface, recording_interface], ["Upload Song", "Record Song"]
)
with gr.Accordion("See all dances", open=False):
gr.Markdown("\n".join(f"- {dance}" for dance in all_dances))
return app
if __name__ == "__main__":
demo().launch()
|