clear 
clc
tic
outputFolder = fullfile('E:\COVID-19\Datasets');
rootFolder = fullfile(outputFolder, 'covidct');
categories = {'CT_COVID', 'CT_NonCOVID'};
imds = imageDatastore(fullfile(rootFolder, categories), 'IncludeSubfolders',true, ...
    'LabelSource', 'foldernames');
[imdsTrain,imdsTest] = splitEachLabel(imds,0.7);

%% Display some sample images.
%%
numTrainImages = numel(imdsTrain.Labels);
idx = randperm(numTrainImages,16);
figure
for i = 1:16
    subplot(4,4,i)
    I = readimage(imdsTrain,idx(i));
    imshow(I)
end
%% Load Pretrained Network

net = resnet18;

net.Layers(1);
inputSize = net.Layers(1).InputSize;
%% Replace Final Layers

if isa(net,'SeriesNetwork') 
  lgraph = layerGraph(net.Layers); 
else
  lgraph = layerGraph(net);
end 
%% 

[learnableLayer,classLayer] = findLayersToReplace(lgraph);
[learnableLayer,classLayer] 
%% 

g = imdsTrain.Labels;
U=unique(g);
numClasses = numel (U);

if isa(learnableLayer,'nnet.cnn.layer.FullyConnectedLayer')
    newLearnableLayer = fullyConnectedLayer(numClasses, ...
        'Name','new_fc', ...
        'WeightLearnRateFactor',20, ...
        'BiasLearnRateFactor',20);
    
elseif isa(learnableLayer,'nnet.cnn.layer.Convolution2DLayer')
    newLearnableLayer = convolution2dLayer(1,numClasses, ...
        'Name','new_conv', ...
        'WeightLearnRateFactor',20, ...
        'BiasLearnRateFactor',20);
end

lgraph = replaceLayer(lgraph,learnableLayer.Name,newLearnableLayer);

newClassLayer = classificationLayer('Name','new_classoutput');
lgraph = replaceLayer(lgraph,classLayer.Name,newClassLayer);

layers = lgraph.Layers;
connections = lgraph.Connections;

layers(1:10) = freezeWeights(layers(1:10));
lgraph = createLgraphUsingConnections(layers,connections);
%% Train Network

pixelRange = [-30 30];
scaleRange = [0.9 1.1];


imageAugmenter = imageDataAugmenter( ...
    'RandXReflection',true, ...
    'RandXTranslation',pixelRange, ...
    'RandYTranslation',pixelRange, ...
    'RandXScale',scaleRange, ...
    'RandYScale',scaleRange); 
augimdsTrain = augmentedImageDatastore(inputSize(1:2),imdsTrain, ...
    'DataAugmentation',imageAugmenter,'ColorPreprocessing','gray2rgb');

augimdsTest = augmentedImageDatastore(inputSize(1:2),imdsTest,'ColorPreprocessing','gray2rgb');

%% 

miniBatchSize = 4;
maxEpochs = 20;
numIterationsPerEpoch = floor(numel(imdsTrain.Labels)/miniBatchSize);
options = trainingOptions('sgdm', ...
    'MiniBatchSize',miniBatchSize, ...
    'MaxEpochs',maxEpochs, ...
    'L2Regularization',0.0005,...
    'InitialLearnRate',1e-4,...
    'ValidationData',augimdsTest, ...
    'ValidationFrequency',numIterationsPerEpoch, ...
    'Verbose',false, ...
    'Plots','training-progress');

New_net = trainNetwork(augimdsTrain,lgraph,options);

%% Alexnet
%featurelayer = 'fc7'; 
%% Google Net 
%featurelayer='pool5-drop_7x7_s1';
%% Resnet 18
featurelayer='pool5';
%% Shuffle net
%featurelayer='node_200';
%
trainingFeatures = activations(New_net,augimdsTrain,featurelayer,'OutputAs','rows');
testFeatures = activations(New_net,augimdsTest,featurelayer,'OutputAs','rows');

%% Extract the class labels from the training and test data.

trainingLabels = imdsTrain.Labels;
testLabels = imdsTest.Labels;
%% Classify Validation Images
% Classify the validation images using the fine-tuned network, and calculate 
% the classification accuracy.

[YPred,probs] = classify(New_net,augimdsTest);
accuracy = mean(YPred == imdsTest.Labels)
%% Calculate the confusion Matrix

confMat = confusionmat(imdsTest.Labels, YPred)

confMat = bsxfun(@rdivide,confMat,sum(confMat,2))   %Convert confusion matrix into percentage form
figure,
plotconfusion(imdsTest.Labels, YPred)

mean(diag(confMat)) % accuracy