%**************************************************************************
% This is a function to train ResNet50 model, iwith 3-channel input.
% Notes:
%   (1)training dataset must be put under folder named ./dataset/train/...
%   (2)testing dataset must be put under folder named ./dataset/test/...
%
% Created By: Kaichang CHENG, Aug 14,2019
%**************************************************************************
%clear;
%close all;
date_now=datestr(now,'yyyy-mm-dd_HH-MM-SS');
net = resnet50; %densenet201
lgraph = layerGraph(net);

% 训练集
trainPath='.\dataset_SouthChinaSea_V3\train';
if(length(dir(trainPath))>2)
imdsTrain = imageDatastore(trainPath,...
    'includeSubfolders',true,...
    'labelsource','foldernames','ReadFcn',@equalScaling224);
else
    disp('No training files.');
    return;
end
T = countEachLabel(imdsTrain);
disp(T);
%测试集
testPath='.\dataset_SouthChinaSea_V3\test';
if(length(dir(testPath))>2)
imdsTest = imageDatastore(testPath,...
    'includeSubfolders',true,...
    'labelsource','foldernames','ReadFcn',@equalScaling224);
else
    disp('No testing files.');
    return;
end
T_test = countEachLabel(imdsTest);
disp(T_test);


lgraph = removeLayers(lgraph, {'fc1000', 'fc1000_softmax','ClassificationLayer_fc1000' });
%lgraph = removeLayers(lgraph, {'input_1','fc1000', 'fc1000_softmax','ClassificationLayer_fc1000' });

numClasses = numel(categories(imdsTrain.Labels));
%inputLayer=imageInputLayer([224 224 1],'Name','input_1''Normalization','zerocenter');
newLayers = [
    fullyConnectedLayer(numClasses,'Name','fc','WeightLearnRateFactor',10,'BiasLearnRateFactor',10)
    softmaxLayer('Name','softmax')
    classificationLayer('Name','classoutput')];
lgraph = addLayers(lgraph,newLayers);
%lgraph = addLayers(lgraph,inputLayer);
lgraph = connectLayers(lgraph,'avg_pool','fc');
%lgraph = connectLayers(lgraph,'input_1','conv1');

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

lgraph = createLgraphUsingConnections(layers,connections);

options = trainingOptions('sgdm',...
    'ValidationData',imdsTest,...
    'ValidationFrequency' ,10,...
    'ValidationPatience',5, ...
    'MiniBatchSize',64,...
    'Maxepochs',5,...%最大迭代轮数，1个epoch等于使用训练集中的全部样本训练一次；1个iteration等于使用batchsize个样本训练一次；
    'InitialLearnRate',0.001,...
    'Shuffle','every-epoch',...
    'Plots','training-progress');


%从训练集划分出验证集
[imds1,imds2] = splitEachLabel(imdsTrain,0.8);
disp('Start training...');
network = trainNetwork(imds1,lgraph,options);

%% predict
tic
predictLabels = classify(network,imdsTest) ;
toc

testLabels = imdsTest.Labels;
%对测试数据标签进行one-hot编码
ttest = dummyvar(double(testLabels))' ;
%对测试数据预测标签进行one-hot编码
tpredictions = dummyvar(double(categorical(predictLabels)))';
%对混淆矩阵做图
plotconfusion(ttest,tpredictions);

accuracy = mean(predictLabels == testLabels);
disp(['Average Accuracy:',num2str(accuracy)]); 
modelName=['network_Resnet50_SouthChinaSea_equalScaling_', date_now, '.mat'];
save(modelName, 'network') ;%保存模型
msgbox('Training complete！','Caution!','warn');