% This program will train yamnet to classify different 
% vehicle starting faults
% 8/8/2024
% Coded in MATLAB R2024b
% Place file under MATLAB folder NIS/Audiofaults
% Note: 
% Need minimum of 6 input files

%Remove previous setting
clear
close all 

%Location for input fault audio files
dataFolder = "./NIS/Audiofaults";

%Setup audio storage resources
ads = audioDatastore(dataFolder,IncludeSubfolders=true,LabelSource="foldernames");

%Setup training and validation partition as 80/20
[adsTrain,adsValidation] = splitEachLabel(ads,0.8,0.2);

[x,fileInfo] = read(adsTrain);
fs = fileInfo.SampleRate;
reset(adsTrain)

%Show time graph of fault audio file
figure
t = (0:size(x,1)-1)/fs;
plot(t,x)
xlabel("Time (s)")
title("State = " + string(fileInfo.Label))
axis tight

%Set training variables
emptyLabelVector = adsTrain.Labels;
emptyLabelVector(:) = [];

trainFeatures = [];
trainLabels = emptyLabelVector;
while hasdata(adsTrain)
    [audioIn,fileInfo] = read(adsTrain);
    features = yamnetPreprocess(audioIn,fileInfo.SampleRate);
    numSpectrums = size(features,4);
    trainFeatures = cat(4,trainFeatures,features);
    trainLabels = cat(2,trainLabels,repmat(fileInfo.Label,1,numSpectrums));
end

validationFeatures = [];
validationLabels = emptyLabelVector;
while hasdata(adsValidation)
    [audioIn,fileInfo] = read(adsValidation);
    features = yamnetPreprocess(audioIn,fileInfo.SampleRate);
    numSpectrums = size(features,4);
    validationFeatures = cat(4,validationFeatures,features);
    validationLabels = cat(2,validationLabels,repmat(fileInfo.Label,1,numSpectrums));
end

classNamesaudio = unique(adsTrain.Labels);
numClasses = numel(classNamesaudio);

net = audioPretrainedNetwork("yamnet",NumClasses=numClasses);

miniBatchSize = 2;
validationFrequency = floor(numel(trainLabels)/miniBatchSize);
options = trainingOptions('adam', ...
    InitialLearnRate=3e-4, ...
    MaxEpochs=2, ...
    MiniBatchSize=miniBatchSize, ...
    Shuffle="every-epoch", ...
    Plots="training-progress", ...
    Metrics="accuracy", ...
    Verbose=false, ...
    ValidationData={single(validationFeatures),validationLabels'}, ...
    ValidationFrequency=validationFrequency);

%Train audio classification tool
AudFN = trainnet(trainFeatures,trainLabels',net,"crossentropy",options);

%Save trained audio tool 
save AudFN.mat