%% Setup of CNN based classifier 

% Fist option direct input to 2D CNN

close all

rootFolder = fullfile('C:\Users\tripats\Documents\Biomedical Signal Analysis\Grad Single Project\ECG');

categories = {'SinusTachycardia',...

    'SinusRhythm','SinusBradycardia',...

     'AtrialFibrilation',};

% categories = {'SupraventricularTachycardia','SinusTachycardia',...

%     'SinusRhythm','SinusIrregularity','SinusBradycardia',...

%     'SinusAtriumtoAtrialWanderingRhythm', 'AtrioventricularReentrantTachycardia',...

%     'AtrioventricularNodeReentrantTachycardia','AtrialTachycardia',...

%      'AtrialFlutter','AtrialFibrilation',};

imds = imageDatastore(fullfile(rootFolder,categories),'LabelSource','foldernames');

imds.ReadFcn = @dsresize;

label_count = countEachLabel(imds)

minCount = min(label_count{:,2})

imds = splitEachLabel(imds, minCount, "randomized")

net = alexnet();

[trainingData,testData] = splitEachLabel(imds, 0.2, 'randomized');

layersTransfer = net.Layers(2:end-3);

layers = [

    imageInputLayer([256 256 3]);

    layersTransfer

    fullyConnectedLayer(numel(categories))

    softmaxLayer

    classificationLayer];

plot(layerGraph(layers));

options = trainingOptions("sgdm",...

    "ExecutionEnvironment","parallel",...

    "InitialLearnRate",1e-3,...

    "MaxEpochs", 40,...

    "Shuffle","every-epoch",...

    "Plots","training-progress",...

    "ValidationData",testData,...

    "ValidationFrequency",5);

[net, traininfo] = trainNetwork(trainingData, layers, options);

function data = dsresize(filename)

    data = imread(filename);

    data = data(:,:,min(1:3, end));

    data = imresize(data, [256 256]);

end