Efficient Retrieval of Music Recordings Using Graph-Based Index Structures

This notebook accompanies the following paper:

• Frank Zalkow, Julian Brandner, and Meinard Müller. Efficient retrieval of music recordings using graph-based index structures. Signals, 2(2):336–352, 2021. doi: 10.3390/signals2020021. Paper Bibtex Website

The paper compares indexing strategies based on $K$-d trees and HNSW graphs in a cross-modal music retrieval application. See more details in the paper. The notebook shows how to load the features for a provided example dataset, generate shingles, reduce the shingle dimensionality using PCA, construct index structures ($K$-d trees and HNSW graphs), and search for the nearest shingles to a given query in the database.

We start by importing some Python packages.

import os
import timeit

import numpy as np
import scipy
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.neighbors import NearestNeighbors, KDTree
import IPython.display as ipd
import hnswlib
from matplotlib import pyplot as plt

from common import compute_features, generate_shingles

Parameters

We now set various parameters used throughout this notebook. First, we decide whether we use dimensionality reduction (binary variable dimensionality_reduction). Second, we fix the used dimensionality $K$ (K, changed to 240 later if no dimensionality reduction is applied). Third, we set the number $\nu$ of neighbors to search for (nu).

Then, we have three more parameters for the HNSW graph, i.e., number $\nu'$ of neighbor candidates during the search (nu_prime), the minimum number $M$ of edges for each node in the graph (M), and the number $M'$ of edge candidates during the construction of the graph (M_prime). For details about the meaning of these parameters, we refer to our paper.

dimensionality_reduction = True
K = 30
nu = 10

nu_prime = 100
M = 5
M_prime = 100

Load Dataset

We next load the metadata and the CENS features for our dataset. Then, we generate shingles (as vectors of 240 dimensions) from the feature sequences and an index that specifies the original document for each shingle.

metadata = pd.read_csv(os.path.join('data', 'metadata.csv'), sep=';')

shingles = []
index = []
 
for _, (cur_idx, cur_composer, cur_work, cur_version, cur_fn) in metadata.iterrows():
    cur_fn_npz = os.path.join('data', 'features_cens', cur_fn)
    cur_X = np.load(cur_fn_npz)['f_chroma']
    cur_shingles = generate_shingles(cur_X)
    
    shingles.append(cur_shingles)
    index.append(np.repeat(cur_idx, cur_shingles.shape[0]))
    
shingles = np.concatenate(shingles, axis=0)
index = np.concatenate(index)

db_size = shingles.shape[0]

To give an overview of the dataset, we create a table with the dataset's musical works, showing the number of versions for each work.

works_with_count = metadata[['ComposerID', 'WorkID']].value_counts(sort=False).to_frame()
works_with_count = pd.DataFrame(works_with_count.to_records())
works_with_count.columns = ['ComposerID', 'WorkID', '# Versions']

ipd.display(ipd.HTML(works_with_count.to_html(index=False)))

ComposerID	WorkID	# Versions
Beethoven	Op055-01	4
Beethoven	Op055-02	4
Beethoven	Op055-03	4
Beethoven	Op055-04	4
Chopin	Op007No1	53
Chopin	Op024No1	61
Chopin	Op033No2	67
Chopin	Op033No3	50
Chopin	Op068No4	62
Vivaldi	RV269-01	7
Vivaldi	RV269-02	7
Vivaldi	RV269-03	7

To give a concrete idea about the dataset content, we now select a random shingle from the dataset, show its corresponding metadata, and visualize the shingle.

example_idx = np.random.randint(0, shingles.shape[0])
example_item = shingles[example_idx, :]

ipd.display(ipd.HTML(metadata[metadata['Index'] == index[example_idx]].to_html()))

plt.figure(figsize=(5, 2.5))
plt.imshow(example_item.reshape(12, 20), aspect='auto', origin='lower', cmap='gray_r', vmin=0, vmax=1)
plt.yticks([0, 2, 4, 5, 7, 9, 11], 'CDEFGAB')
plt.ylabel('Chroma')
plt.xlabel('Time (seconds)')
plt.colorbar()
plt.tight_layout()

	Index	ComposerID	WorkID	VersionID	Filename
309	309	Vivaldi	RV269-01	Abbado	Vivaldi_RV269-01_Abbado.npz

Dimensionality Reduction

If required, we perform dimensionality reduction to the shingle using PCA.

if dimensionality_reduction:
    pca = PCA(n_components=K)
    pca.fit(shingles)
    items = pca.transform(shingles)
else:
    items = shingles
    K = shingles.shape[1]

Query

In this notebook, we use a query from the database to demonstrate the usage of the index structures. You may also use another query by computing the features for an audio file. To do this, we provide the functions common.compute_features and common.generate_shingles.

query_idx = np.random.randint(0, items.shape[0])
query_item = items[query_idx, :]

print('Query')
ipd.display(ipd.HTML(metadata[metadata['Index'] == index[query_idx]].to_html()))

# to use another query, you can compute a query item from an audio file
# import librosa
# fn_audio = '...'  # please provide the path to an audio file
# X_cens = compute_features(fn_audio)
# query_shingle = generate_shingles(X_cens)[0, :]  # using the first shingle from the file
# query_item = pca.transform(query_shingle.reshape(1, -1))[0] if dimensionality_reduction else query_shingle

Query

	Index	ComposerID	WorkID	VersionID	Filename
46	46	Chopin	Op007No1	Kushner-1990_pid9118-05	Chopin_Op007No1_Kushner-1990_pid9118-05.npz

Exhaustive Nearest Neighbor Search

Now, we search for the $\nu$ nearest shingles in our dataset using an exhaustive search. Here, we compute a matrix of pairwise distances between the query and all database items and select the items corresponding to the $\nu$ smallest distances.

def exhaustive_search(database, queries, n_neighbors):
    C = scipy.spatial.distance.cdist(database, queries, metric='euclidean')
    neighbor_idx = np.argsort(C, axis=0)[:n_neighbors, :]
    neighbor_dist = C[np.unravel_index(neighbor_idx, C.shape)]
    return neighbor_idx, neighbor_dist
    
neighbor_idx, neighbor_dist = exhaustive_search(items, query_item.reshape(1, -1), nu)
neighbor_dist = np.square(neighbor_dist)  # show squared Euclidean distance

We next show the results of the search using the function show_result_table.

def show_result_table(neighbor_idx, neighbor_dist, index=index, metadata=metadata):
    neighbor_idx = neighbor_idx.ravel()
    neighbor_dist = neighbor_dist.ravel()
    
    table = []
    for i, (cur_item_idx, cur_dist) in enumerate(zip(neighbor_idx, neighbor_dist)):
        cur_version_idx = index[cur_item_idx]
        cur_meta_row = metadata[metadata['Index'] == cur_version_idx].iloc[0]
        
        table.append([f'{i + 1:02d}', cur_meta_row['ComposerID'], cur_meta_row['WorkID'],
                      cur_meta_row['VersionID'], cur_item_idx, cur_dist])
    
    cols = ['Match', 'ComposerID', 'WorkID', 'VersionID', 'ShingleIdx', 'Distance']
    ipd.display(ipd.HTML(pd.DataFrame(table, columns=cols).to_html(index=False)))
    
print('Results of exhaustive search')
show_result_table(neighbor_idx, neighbor_dist)

Results of exhaustive search

Match	ComposerID	WorkID	VersionID	ShingleIdx	Distance
01	Chopin	Op007No1	Kushner-1990_pid9118-05	15251	0.000000
02	Chopin	Op007No1	Kushner-1990_pid9118-05	15226	0.342614
03	Chopin	Op007No1	Kushner-1990_pid9118-05	15252	0.393929
04	Chopin	Op007No1	Kushner-1990_pid9118-05	15250	0.408689
05	Chopin	Op007No1	Kushner-1990_pid9118-05	15227	0.488261
06	Chopin	Op007No1	Ohlsson-1999_pid9153-05	16364	0.587715
07	Chopin	Op007No1	Czerny-Stefanska-1949_pid9086-05	13113	0.635510
08	Chopin	Op007No1	Czerny-Stefanska-1949_pid9086-05	13141	0.635921
09	Chopin	Op007No1	Uninsky-1959_pid9061-05	17532	0.656634
10	Chopin	Op007No1	Uninsky-1959_pid9061-05	17557	0.665450

Nearest Neighbor Search with $K$-d tree

Next, we construct a $K$-d tree for our dataset and search for the $\nu$ nearest shingles using this tree. We then show the results as a table.

neighbors = NearestNeighbors(n_neighbors=nu, metric='euclidean')
neighbors.fit(items)

assert isinstance(neighbors._tree, KDTree)

neighbor_dist, neighbor_idx = neighbors.kneighbors(query_item.reshape(1, -1), nu)
neighbor_dist = np.square(neighbor_dist)  # show squared Euclidean distance

print('Results of search with K-d tree')
show_result_table(neighbor_idx, neighbor_dist)

Results of search with K-d tree

Match	ComposerID	WorkID	VersionID	ShingleIdx	Distance
01	Chopin	Op007No1	Kushner-1990_pid9118-05	15251	0.000000
02	Chopin	Op007No1	Kushner-1990_pid9118-05	15226	0.342614
03	Chopin	Op007No1	Kushner-1990_pid9118-05	15252	0.393929
04	Chopin	Op007No1	Kushner-1990_pid9118-05	15250	0.408689
05	Chopin	Op007No1	Kushner-1990_pid9118-05	15227	0.488261
06	Chopin	Op007No1	Ohlsson-1999_pid9153-05	16364	0.587715
07	Chopin	Op007No1	Czerny-Stefanska-1949_pid9086-05	13113	0.635510
08	Chopin	Op007No1	Czerny-Stefanska-1949_pid9086-05	13141	0.635921
09	Chopin	Op007No1	Uninsky-1959_pid9061-05	17532	0.656634
10	Chopin	Op007No1	Uninsky-1959_pid9061-05	17557	0.665450

Nearest Neighbor Search with HNSW Graph

Finally, we construct an HNSW graph for our dataset and search for the $\nu$ nearest shingles using this graph. We then show the results as a table.

graph = hnswlib.Index(space='l2', dim=K)
graph.init_index(max_elements=items.shape[0], ef_construction=M_prime, M=M)
graph.set_ef(nu_prime)
graph.add_items(items, np.arange(items.shape[0]))

neighbor_idx, neighbor_dist = graph.knn_query(query_item.reshape(1, -1), k=nu)

print('Results of search with HNSW graph')
show_result_table(neighbor_idx, neighbor_dist)

Results of search with HNSW graph

Match	ComposerID	WorkID	VersionID	ShingleIdx	Distance
01	Chopin	Op007No1	Kushner-1990_pid9118-05	15251	0.000000
02	Chopin	Op007No1	Kushner-1990_pid9118-05	15226	0.342614
03	Chopin	Op007No1	Kushner-1990_pid9118-05	15252	0.393929
04	Chopin	Op007No1	Kushner-1990_pid9118-05	15250	0.408689
05	Chopin	Op007No1	Kushner-1990_pid9118-05	15227	0.488261
06	Chopin	Op007No1	Ohlsson-1999_pid9153-05	16364	0.587715
07	Chopin	Op007No1	Czerny-Stefanska-1949_pid9086-05	13113	0.635511
08	Chopin	Op007No1	Czerny-Stefanska-1949_pid9086-05	13141	0.635921
09	Chopin	Op007No1	Uninsky-1959_pid9061-05	17532	0.656634
10	Chopin	Op007No1	Uninsky-1959_pid9061-05	17557	0.665450

Retrieval Runtime

The aim of using index structures in our music application is to increase the efficiency of the retrieval runtime. In the next code cell, we compare the runtimes of the different search strategies. First, we define functions with a consistent interface for the exhaustive search, the $K$-d tree search, and the HNSW-based search. Then we measure the runtimes of these functions using the Python module timeit.

def kd(neighbors, queries, n_neighbors):
    neighbor_dist, neighbor_idx = neighbors.kneighbors(queries, nu)
    return neighbor_idx, neighbor_dist

def hnsw(graph, queries, n_neighbors):
    neighbor_idx, neighbor_dist = graph.knn_query(queries, k=nu)
    return neighbor_idx, neighbor_dist
    
# measure runtimes
times_fullsearch = timeit.repeat('exhaustive_search(items, query_item.reshape(1, -1), nu)',
                                 globals=globals(), number=1, repeat=100)
times_kd = timeit.repeat('kd(neighbors, query_item.reshape(1, -1), nu)',
                         globals=globals(), number=1, repeat=100)
times_hnsw = timeit.repeat('hnsw(graph, query_item.reshape(1, -1), nu)',
                           globals=globals(), number=1, repeat=100)

# convert from second to milliseconds
times_fullsearch = np.array(times_fullsearch) * 1000
times_kd = np.array(times_kd) * 1000
times_hnsw = np.array(times_hnsw) * 1000

# report results
print(f'Full Search: {np.mean(times_fullsearch):.3f} ms +/- {np.std(times_fullsearch):.1f} ms')
print(f'KD:          {np.mean(times_kd):.3f} ms +/- {np.std(times_kd):.1f} ms')
print(f'HNSW:        {np.mean(times_hnsw):.3f} ms +/- {np.std(times_hnsw):.1f} ms')

Full Search: 4.942 ms +/- 1.0 ms
KD:          0.797 ms +/- 0.5 ms
HNSW:        0.023 ms +/- 0.0 ms

Acknowledgment: This notebook was created by Frank Zalkow.