%% Transfer Learning Using VGGish
% In this example, you transfer the learning in the VGGish regression model to an audio classification task.
%
% Download and unzip the environmental sound classification data set. This data set consists of recordings labeled as one of 10 different audio sound classes (ESC-10).

url = 'http://ssd.mathworks.com/supportfiles/audio/ESC-10.zip';
downloadFolder = fullfile(tempdir,'ESC-10');
datasetLocation = tempdir;

if ~exist(fullfile(tempdir,'ESC-10'),'dir')
    loc = websave(downloadFolder,url);
    unzip(loc,fullfile(tempdir,'ESC-10'))
end

%
% Create an audioDatastore object to manage the data and split it into train and validation sets. Call countEachLabel to display the distribution of sound classes and the number of unique labels.

ads = audioDatastore(downloadFolder,'IncludeSubfolders',true,'LabelSource','foldernames');
labelTable = countEachLabel(ads)

%
%  Determine the total number of classes.

numClasses = size(labelTable,1);

%
% Call splitEachLabel to split the data set into training and validation sets. Inspect the distribution of labels in the training and validation sets.

[adsTrain, adsValidation] = splitEachLabel(ads,0.8);

countEachLabel(adsTrain)
countEachLabel(adsValidation)

%
% The VGGish network expects audio to be preprocessed into log mel spectrograms. The supporting function vggishPreprocess takes an audioDatastore object and the overlap percentage between log mel spectrograms as input, and returns matrices of predictors and responses suitable as input to the VGGish network.

overlapPercentage = 75;

[trainFeatures,trainLabels] = vggishPreprocess(adsTrain,overlapPercentage);
[validationFeatures,validationLabels,segmentsPerFile] = vggishPreprocess(adsValidation,overlapPercentage);

%
% Load the VGGish model and convert it to a layerGraph object.

net = vggish;

lgraph = layerGraph(net.Layers);

%
% Use removeLayers to remove the final regression output layer from the graph. After you remove the regression layer, the new final layer of the graph is a ReLU layer named 'EmbeddingBatch'.

lgraph = removeLayers(lgraph,'regressionoutput');
lgraph.Layers(end)

%
% Use addLayers to add a fullyConnectedLayer, a softmaxLayer, and a classificationLayer to the graph.

lgraph = addLayers(lgraph,fullyConnectedLayer(numClasses,'Name','FCFinal'));
lgraph = addLayers(lgraph,softmaxLayer('Name','softmax'));
lgraph = addLayers(lgraph,classificationLayer('Name','classOut'));

%
% Use connectLayers to append the fully connected, softmax, and classification layers to the layer graph.

lgraph = connectLayers(lgraph,'EmbeddingBatch','FCFinal');
lgraph = connectLayers(lgraph,'FCFinal','softmax');
lgraph = connectLayers(lgraph,'softmax','classOut');

%
% To define training options, use trainingOptions.

miniBatchSize = 128;
options = trainingOptions('adam', ...
    'MaxEpochs',5, ...
    'MiniBatchSize',miniBatchSize, ...
    'Shuffle','every-epoch', ...
    'ValidationData',{validationFeatures,validationLabels}, ...
    'ValidationFrequency',50, ...
    'LearnRateSchedule','piecewise', ...
    'LearnRateDropFactor',0.5, ...
    'LearnRateDropPeriod',2);

%
% To train the network, use trainNetwork.

[trainedNet, netInfo] = trainNetwork(trainFeatures,trainLabels,lgraph,options);

%
% Each audio file was split into several segments to feed into the VGGish network. Combine the predictions for each file in the validation set using a majority-rule decision.

validationPredictions = classify(trainedNet,validationFeatures);

idx = 1;
validationPredictionsPerFile = categorical;
for ii = 1:numel(adsValidation.Files)
    validationPredictionsPerFile(ii,1) = mode(validationPredictions(idx:idx+segmentsPerFile(ii)-1));
    idx = idx + segmentsPerFile(ii);
end

%
% Use confusionChart to evaluate the performance of the network on the validation set.

figure('Units','normalized','Position',[0.2 0.2 0.5 0.5]);
cm = confusionchart(adsValidation.Labels,validationPredictionsPerFile);
cm.Title = sprintf('Confusion Matrix for Validation Data \nAccuracy = %0.2f %%',mean(validationPredictionsPerFile==adsValidation.Labels)*100);
cm.ColumnSummary = 'column-normalized';
cm.RowSummary = 'row-normalized';

%
%% Supporting Functions

function [predictor,response,segmentsPerFile] = vggishPreprocess(ads,overlap)
% This function is for example purposes only and may be changed or removed
% in a future release.

% Create filter bank
FFTLength = 512;
numBands = 64;
fs0 = 16e3;
filterBank = designAuditoryFilterBank(fs0, ...
    'FrequencyScale','mel', ...
    'FFTLength',FFTLength, ...
    'FrequencyRange',[125 7500], ...
    'NumBands',numBands, ...
    'Normalization','none', ...
    'FilterBankDesignDomain','warped');

% Define STFT parameters
windowLength = 0.025 * fs0;
hopLength = 0.01 * fs0;
win = hann(windowLength,'periodic');

% Define spectrogram segmentation parameters
segmentDuration = 0.96; % seconds
segmentRate = 100; % hertz
segmentLength = segmentDuration*segmentRate; % Number of spectrums per auditory spectrograms
segmentHopDuration = (100-overlap) * segmentDuration / 100; % Duration (s) advanced between auditory spectrograms
segmentHopLength = round(segmentHopDuration * segmentRate); % Number of spectrums advanced between auditory spectrograms

% Preallocate cell arrays for the predictors and responses
numFiles = numel(ads.Files);
predictor = cell(numFiles,1);
response = predictor;
segmentsPerFile = zeros(numFiles,1);

% Extract predictors and responses for each file
for ii = 1:numFiles
    [audioIn,info] = read(ads);

    x = single(resample(audioIn,fs0,info.SampleRate));

    Y = stft(x, ...
        'Window',win, ...
        'OverlapLength',windowLength-hopLength, ...
        'FFTLength',FFTLength, ...
        'FrequencyRange','onesided');
    Y = abs(Y);

    logMelSpectrogram = log(filterBank*Y + single(0.01))';
    
    % Segment log-mel spectrogram
    numHops = floor((size(Y,2)-segmentLength)/segmentHopLength) + 1;
    segmentedLogMelSpectrogram = zeros(segmentLength,numBands,1,numHops);
    for hop = 1:numHops
        segmentedLogMelSpectrogram(:,:,1,hop) = logMelSpectrogram(1+segmentHopLength*(hop-1):segmentLength+segmentHopLength*(hop-1),:);
    end

    predictor{ii} = segmentedLogMelSpectrogram;
    response{ii} = repelem(info.Label,numHops);
    segmentsPerFile(ii) = numHops;
end

% Concatenate predictors and responses into arrays
predictor = cat(4,predictor{:});
response = cat(2,response{:});
end

% Copyright 2020 The MathWorks, Inc.