%%%%10.01.20 train network for classification using matlab R2017b
    %%%Workflow:
        %%% 1. Label data or obtain labeled data
        %%% 2. Create a datastore for original and labeled images
        %%% 3. Partition the datastores in training and test sets
        %%% 4. Import a CNN and modify to be a SegNet
        %%% 5. Train and Evaluate the network
%% Resize images and pixel labels (added 10.13.20)
% path_in_img='/gpfs/summit/scratch/apsendro@colostate.edu/AllTiles'
% path_in_pixel='/gpfs/summit/scratch/apsendro@colostate.edu/MackTiles_corrected/PixelLabelData';
% 
% path_out_img='/gpfs/summit/scratch/apsendro@colostate.edu/AllTilesResized'
% path_out_pixel='/gpfs/summit/scratch/apsendro@colostate.edu/MackTiles_corrected/PixelLabelDataResized';
% 
% for ii=1:10053;
%     %%for the images
%     q=['tile_',num2str(ii),'.tif'];
%     img=imread(fullfile(path_in_img,q));
%     
%     img_new=img(1:240,1:240,:);
%     
%     filename=fullfile(path_out_img,q);
%     imwrite(img_new,filename);
%     
%     clear q img img_new filename
%     
%     %%for the pixel labels
%     p=['Label_',num2str(ii),'.png'];
%     px=imread(fullfile(path_in_pixel,p));
%     
%     px_new=px(1:240,1:240);
%     
%     filename1=fullfile(path_out_pixel,p);
%     imwrite(px_new,filename1);
%     
%     clear p px px_new filename1
% end
% 
% clear ii path_in_img path_in_pixel path_out_img path_out_pixel
%         
        
%% 1&2- load data into datastores
%%%create ImageDatastore and PixelLabelDatastore
path_img='/gpfs/summit/scratch/apsendro@colostate.edu/AllTilesResized'
path_pixel='/gpfs/summit/scratch/apsendro@colostate.edu/MackTiles_corrected/PixelLabelDataResized';
classNames={'Wood';'Background'};
pixelLabelIDs=[1;2];

imds=imageDatastore(path_img);
pxds = pixelLabelDatastore(path_pixel,classNames,pixelLabelIDs);

%%%see distribution of class labels
tbl = countEachLabel(pxds);

    %%%balance classes using class weighting
    imageFreq = tbl.PixelCount ./ tbl.ImagePixelCount;
    classWeights = median(imageFreq) ./ imageFreq
    
%% 3- Partition stores
%%%partition data into training and test set, choose n: proportion of data
%%%used for training while the rest is used for testing
n=0.7;
[imdsTrain, imdsTest, pxdsTrain, pxdsTest] = partitionData(imds,pxds,n);

    %%%find the no. of training and test images
    numTrainingImages = numel(imdsTrain.Files);
    numTestingImages = numel(imdsTest.Files);

% %%%set up training data
% trainingData = pixelLabelImageSource(imdsTrain,pxdsTrain,'OutputSize',[240 240],'OutputSizeMode','centercrop','BackgroundExecution',true);
trainingData = pixelLabelImageSource(imdsTrain,pxdsTrain,'BackgroundExecution',true);

%% 4- Create segnet layers, set hyperparameter options

% load('net_epoch7.mat','net')
% 
imageSize=[240 240 3];
numClasses = 2;
encoderDepth=4;
lgraph=segnetLayers(imageSize,numClasses,encoderDepth);

%%%create new PixelClassificationLayer taking into account class weights
pxLayer = pixelClassificationLayer('Name','Seg-Layer','ClassNames',classNames,'ClassWeights',classWeights);

%%%replace layers in network
lgraph=removeLayers(lgraph,lgraph.Layers(59,1).Name);
lgraph=addLayers(lgraph,pxLayer);
lgraph=connectLayers(lgraph,lgraph.Layers(58,1).Name,lgraph.Layers(59,1).Name);

%%%set options
check_path='/gpfs/summit/scratch/apsendro@colostate.edu/mackTrainedNetwork';

options = trainingOptions('sgdm',...
    'Momentum',0.9,...
    'InitialLearnRate',0.03, ...
    'LearnRateSchedule','piecewise',... 
    'LearnRateDropPeriod',10,...
    'LearnRateDropFactor',0.1,...
    'L2Regularization',0.0001,...
    'MaxEpochs',2,...
    'MiniBatchSize',128,...
    'VerboseFrequency',50,...
    'Shuffle','every-epoch',...
    'CheckpointPath',check_path,...
    'SequenceLength','longest',...
    'ExecutionEnvironment','cpu',...
    'OutputFcn',@(info)stopIfAccuracyNotImproving(info,5));

%% 5- train and evaluate network

% net2 = trainNetwork(trainingData,layerGraph(net),options);
% 
% save(fullfile(check_path,'network2.mat'),'net2');

%%evaluate trained network on test set
load(fullfile(check_path,'network2.mat'));

pxdsResults=semanticseg(imdsTest,net2,'MiniBatchSize',256);

metrics=evaluateSemanticSegmentation(pxdsResults,pxdsTest);

save(fullfile(check_path,'AllResults.mat'));