File size: 2,266 Bytes
c32e5ef | 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 | """LAION-CLAP's music checkpoint on the same benchmark as song_describer.py, for comparison.
uv run clap_eval.py
"""
import argparse
from pathlib import Path
import laion_clap
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from song_describer import clips, load_captions, load_tracks, print_results
SR = 48000
CHECKPOINT = "music_audioset_epoch_15_esc_90.14.pt"
def normalize(vectors):
return vectors / np.linalg.norm(vectors, axis=-1, keepdims=True)
@torch.inference_mode()
def embed(model, wave):
size = 10 * SR
chunks = np.stack([wave[start:start + size] for start in range(0, len(wave) - size + 1, size)])
# CLAP was trained on audio that went through int16.
chunks = (np.clip(chunks, -1, 1) * 32767).astype(np.int16).astype(np.float32) / 32767
return normalize(model.get_audio_embedding_from_data(chunks).mean(0))
@torch.inference_mode()
def embed_text(model, captions):
return normalize(np.concatenate([model.get_text_embedding(captions[start:start + 64]) for start in range(0, len(captions), 64)]))
def main():
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--data", type=Path, default=Path("data/song_describer"))
parser.add_argument("--device", default="cuda", help="cuda, mps or cpu")
args = parser.parse_args()
captions = load_captions(args.data)
model = laion_clap.CLAP_Module(enable_fusion=False, amodel="HTSAT-base", device=args.device)
model.load_ckpt(hf_hub_download("lukewys/laion_clap", CHECKPOINT), verbose=False)
tracks, middle, average = [], [], []
for track, wave in load_tracks(args.data, SR):
three = [embed(model, clip) for clip in clips(wave, SR)]
tracks.append(track)
middle.append(three[1])
average.append(normalize(np.mean(three, axis=0)))
order = np.argsort(tracks)
views = {"one 10 s clip": middle, "three 10 s clips": average}
views = {name: np.stack(audio)[order] for name, audio in views.items()}
text = embed_text(model, captions.text.tolist())
print_results("CLAP music", text, views, captions["corpus-id"].to_numpy(), np.array(tracks)[order])
if __name__ == "__main__":
main()
|