CLAPを使ったお手軽な音声検索を実装してみた

概要

Contrastive Language-Audio Pretraining (CLAP)はCLIPの音声版です。
https://arxiv.org/abs/2211.06687

CLIPを使った画像検索は記事がいくつかありますが、CLAPを使った音声検索は見当たらなかったので実装してみました。今回は「音声」で「近い音声」を検索します。

使用するモデル

今回はこちらのLAION-AIのCLAPの実装を利用します。CLAPはMSの実装もあるようですが、LAION-AIのほうがよく使われているようです。
https://github.com/LAION-AI/CLAP

テスト用にAudioSetデータセットをダウンロード

今回はAudioSetデータセットの一部をダウンロードしてテストデータとして使います。

AudioSetのダウンロードにはaudioset-downloadモジュールを使います。
https://github.com/MorenoLaQuatra/audioset-download

apt installでffmpegをインストールした後で、pipでaudioset_downloadをインストールします。

sudo apt install ffmpeg
pip install audioset_download

download_type='eval'として、評価用のデータセットをダウンロードします。CLAPの事前学習済みモデルはAudioSetで学習しているのでtrainのほうは使わないほうがよさそうです。

from audioset_download import Downloader
d = Downloader(root_path='eval_wav', labels=None, n_jobs=16, download_type='eval', copy_and_replicate=False)
d.download(format = 'wav')

こちらは時間がかかるので、動かしたまま次に行きます。

CLAPの事前学習済みモデルのロード

ここは公式サイトに従います。デフォルトではAudioSetデータセットで学習した事前学習済みモデルがロードされます。楽曲や会話の音声で学習した事前学習モデルも公開されています。LAION-AI/CLAPのGitHubを参照してください。

import numpy as np
import librosa
import torch
import laion_clap
import glob

# quantization
def int16_to_float32(x):
    return (x / 32767.0).astype(np.float32)

def float32_to_int16(x):
    x = np.clip(x, a_min=-1., a_max=1.)
    return (x * 32767.).astype(np.int16)

model = laion_clap.CLAP_Module(enable_fusion=False)
model.load_ckpt() # download the default pretrained checkpoint.

テストデータを読み込んでベクトル変換

ファイルを一つずつ読み込んで埋め込みベクトルに変換します。get_audio_embedding_from_filelist()関数には複数のwavファイルのリストを渡せるのですが、たまにエラーになってしまうので一つずつ変換しています。

import os
from tqdm import tqdm

input_dir = 'eval_wav/**/*.wav'
file_list = [p for p in glob.glob(input_dir, recursive=True) if os.path.isfile(p)]

print('file num:', len(file_list))
#file_list = file_list[:100]

def get_class_name(file):
    dir_name = os.path.dirname(file)
    return dir_name.split('/')[1]

info_list = []
for file in tqdm(file_list):
    try:
        audio_embed = model.get_audio_embedding_from_filelist(x=[file], use_tensor=False)        
            
        info = {}
        info['file'] = file
        info['class'] = get_class_name(file)
        info['embed'] = audio_embed[0]
        info_list.append(info)
            
    except Exception as e:
        print('error:', e)
        continue

print(info_list[0])

埋め込みベクトルで近い音声を検索

これですべての音声の埋め込みベクトルが求まりました。適当に一つ取り出して、それと近い順にソートしてみます。

import numpy as np
import random

# 検索対象のインデックスをランダムに決める
find_index = random.randint(0, len(info_list)-1)

embed_list = [info['embed'] for info in info_list]
dot_product_list = []

# 検索対象の埋め込みベクトル
find_embed = embed_list[find_index]

# dot productを取ってコサイン類似度を求める
for embed in embed_list:
    dot_product_list.append(np.dot(find_embed, embed))
    
print(dot_product_list[:10])    
    
sorted_index_list = list(np.argsort(dot_product_list))
sorted_index_list.reverse()

print('find index class:', info_list[find_index]['class'])
for i in sorted_index_list[:50]:
    print(info_list[i]['class'], ' : ', dot_product_list[i])

実行結果は以下のようになります。Background musicクラスに近い音声ということで、music系のクラスの音声が上位に来ていることがわかります。

find index class: Background music
Background music  :  1.0
Theme music  :  0.916441
Christmas music  :  0.73565924
Music  :  0.71548736
Music  :  0.7049972
Orchestra  :  0.70190996
Opera  :  0.6827497
Music  :  0.67603403
Background music  :  0.67359114
Traditional music  :  0.67343926
Music  :  0.6623305
Christmas music  :  0.6617918
Traditional music  :  0.6600741
String section  :  0.6595639
String section  :  0.64624965
Background music  :  0.64507854
New-age music  :  0.64478624
Music  :  0.6401868
Organ  :  0.64011776
Music  :  0.6354705

ベクトルに変換後はCLIPと変わらないので、高速化のためにfaissを使った検索を行うことができますが、今回は省略します。