src.eval

  1import numpy as np
  2import pandas as pd
  3import torch
  4from sklearn.metrics import ConfusionMatrixDisplay
  5from torch.utils.data import DataLoader
  6from tqdm import tqdm
  7
  8from src.data import DatasetMusic
  9from src.data import DatasetMusicTest
 10from src.network import AudioNetwork
 11
 12
 13def get_anchor_embeddings(network: torch.nn.Module, anchors_waveform: dict) -> dict:
 14    # Get the embeddings for each anchor
 15    anchors_embedding = {}
 16    for name, waveform in anchors_waveform.items():
 17        # Perform inferencing to get the embedding
 18        with torch.no_grad():
 19            output = network(waveform.unsqueeze(0))
 20            anchors_embedding[name] = output.detach().numpy()
 21
 22    return anchors_embedding
 23
 24
 25def plot_confusion_matrix():
 26    # Variables
 27    model_save_path = "save/audio_network.pkl"
 28
 29    # Create the network
 30    network = AudioNetwork()
 31    network.load_state_dict(torch.load(model_save_path))
 32
 33    # Create the dataloader for validation
 34    dataset_valid = DatasetMusic(train=False, spectrogram=True)
 35    dataloader_valid = DataLoader(dataset_valid, batch_size=1)
 36
 37    # Get the anchor waveforms
 38    anchors_waveforms = dataset_valid.get_anchors()
 39    anchors_embedding = get_anchor_embeddings(network, anchors_waveforms)
 40
 41    y_true = []
 42    y_pred = []
 43
 44    for sample in tqdm(dataloader_valid, desc="Validating", colour="red"):
 45        # Waveforms
 46        waveform_pos = sample[1]
 47
 48        # Labels
 49        label_pos = sample[4][0]
 50
 51        # Generate prediction
 52        prediction = network(waveform_pos)
 53
 54        # Need to compare the prediction versus all anchors
 55        association = None
 56        max_dist = 1000
 57        for key, val in anchors_embedding.items():
 58            # We are going to use Euclidean distance to figure out what to associate the embedding with
 59            dist = np.linalg.norm(prediction.detach().numpy() - val)
 60            if dist < max_dist:
 61                max_dist = dist
 62                association = key
 63
 64        # Check if dist is out of bounds
 65        if max_dist > 5:
 66            association = "Other"
 67
 68        # Assign confusion matrix values
 69        y_true.append(label_pos)
 70        y_pred.append(association)
 71
 72    ConfusionMatrixDisplay.from_predictions(
 73        y_true,
 74        y_pred,
 75        labels=["Beethoven", "Schubert", "Bach", "Brahms", "Other"],
 76    )
 77
 78    print("Stop here to view the confusion matrix!")
 79
 80
 81def infer_test_data():
 82    # Variables
 83    model_save_path = "save/audio_network.pkl"
 84
 85    # Create the network
 86    network = AudioNetwork()
 87    network.load_state_dict(torch.load(model_save_path))
 88    network.eval()
 89
 90    # Create the dataloader for validation
 91    dataset_valid = DatasetMusic(train=False, spectrogram=True)
 92    dataset_test = DatasetMusicTest(spectrogram=True)
 93
 94    # Setup dataloader
 95    dataloader_test = DataLoader(dataset_test, batch_size=1)
 96
 97    # Get the anchor waveforms
 98    anchors_waveforms = dataset_valid.get_anchors()
 99    anchors_embedding = get_anchor_embeddings(network, anchors_waveforms)
100
101    distances = []
102
103    for sample in tqdm(dataloader_test, desc="Validating", colour="red"):
104        # Waveforms
105        waveform = sample[0]
106        filename = sample[1][0]
107
108        # Generate prediction
109        prediction = network(waveform)
110
111        # Need to compare the prediction versus all anchors
112        min_dist = 1000
113        for key, val in anchors_embedding.items():
114            # We are going to use Euclidean distance to figure out what to associate the embedding with
115            dist = np.linalg.norm(prediction.detach().numpy() - val)
116            if dist < min_dist:
117                min_dist = dist
118
119        # Check if dist is out of bounds
120        distances.append((filename, min_dist))
121
122    # Print out the mininum distances
123    df = pd.DataFrame(distances, columns=["filename", "min_dist"])
124    df = df.sort_values("min_dist").reset_index(drop=True)
125    print(df)
126    print("done!")
127
128
129if __name__ == "__main__":
130    plot_confusion_matrix()
131    # infer_test_data()