%%%%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('network29.mat','net29')

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/mackNetworkParameters/addTrain1';

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

%% 5- train and evaluate network

net30 = trainNetwork(trainingData,layerGraph(net29),options);

save(fullfile(check_path,'network30.mat'),'net30');

%%evaluate trained network on test set

path_px_test='/gpfs/summit/scratch/apsendro@colostate.edu/mackNetworkParameters/addTrain1/pix9';

pxdsResults=semanticseg(imdsTest,net30,'MiniBatchSize',256,'WriteLocation',path_px_test);
% pxdsResults=semanticseg(imdsTest,net21,'MiniBatchSize',256);

metrics=evaluateSemanticSegmentation(pxdsResults,pxdsTest);

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

%% 6- new addition 01.07.21-- checking the quickbird image as a test
path_in_qb='/gpfs/summit/scratch/apsendro@colostate.edu/mackClassify/Images/splitBatch2'
path_out_qb='/gpfs/summit/scratch/apsendro@colostate.edu/mackClassify/Classified/splitBatch2_c'

files=dir(fullfile(path_in_qb,'*.TIF'));

rbg1=[3,2,1];
rbg2=[5,3,2];
step=1024;

for zz=14
    zz

    [img,R] = geotiffread(fullfile(path_in_qb,files(zz).name));
    info = geotiffinfo(fullfile(path_in_qb,files(zz).name));

    key = info.GeoTIFFTags.GeoKeyDirectoryTag;

    s=size(img);
    if s(3)<6;
       img_rbg=img(:,:,rbg1);
    elseif s(3)>7; 
       img_rbg=img(:,:,rbg2);
    end
    clear s 
    
    a=size(img_rbg);
    
    C=uint8(NaN(a(1),a(2)));
        
    n_h = fix(a(1)/step);
    n_w = fix(a(2)/step);
    
    for ii=0:n_w-1
        for jj=0:n_h-1
    
            x=1+(step*jj);
            y=1+(step*ii);
            width=step;
            height=step;
        
            C(x:x+width,y:y+height)=semanticseg(img_rbg(x:x+width,y:y+height,:),net30,'OutputType','double','ExecutionEnvironment','cpu');
            
        end
    end
    clear ii jj x y width height

    C((a(1)-step:a(1)),:)=semanticseg(img_rbg((a(1)-step:a(1)),:,:),net30,'OutputType','double','ExecutionEnvironment','cpu');
    C(:,(a(2)-step):a(2))=semanticseg(img_rbg(:,(a(2)-step):a(2),:),net30,'OutputType','double','ExecutionEnvironment','cpu');
    
    wood=sum(C(:)==1);
    wood_cover=wood*R.CellExtentInWorldX*R.CellExtentInWorldY;
    wood_cover
    
    [blob,name1]=fileparts(files(zz).name);
    namesave=fullfile(path_out_qb,strcat(name1,'_classified_8_30.tif'));
    geotiffwrite(namesave,C,R,'GeoKeyDirectoryTag',key);
    
clear img R info key img_rbg a C n_h n_w name1 namesave wood blob
end

clear files path_in path_out rbg1 rbg2 step zz p
