%% This code train and test semantic segmentation of Capsule endoscopy images
% required MATLAB 2018
% Developer: Tonmoy Ghosh (tghosh@crimson.ua.edu)

% Load Images
% Use |imageDatastore| to load images. The |imageDatastore| enables you 
% to efficiently load a large collection of images on disk.
%%
clear; clc; close all;
%imgDir = fullfile('bleeding images')
imgDir = '/Users/tonmoyghosh/OneDrive - The University of Alabama/Paper with Code/Semantic Segmentation Based Bleeding Zone Detection/Dataset/bleeding';
imds = imageDatastore(imgDir);
%% 
% Display one of the images.

I = readimage(imds, 1);
I = histeq(I);
figure
imshow(I)
%% Load Pixel-Labeled Images
classes = [
    "Bleeding"
    "Non_Bleeding"
    "Background"
    ];
labelIDs = PixelLabelIDs();
%% 
% Use the classes and label IDs to create the |pixelLabelDatastore|:

%labelDir = fullfile('labels');
labelDir = '/Users/tonmoyghosh/OneDrive - The University of Alabama/Paper with Code/Semantic Segmentation Based Bleeding Zone Detection/Dataset/labels';
pxds = pixelLabelDatastore(labelDir,classes,labelIDs);
% Read and display one of the pixel-labeled images by overlaying it on top 
% of an image.

C = readimage(pxds, 1);


cmap = CEColorMap;
B = labeloverlay(I,C,'ColorMap',cmap);

figure
imshow(B)
pixelLabelColorbar(cmap,classes);

%%
%analize the data statistics
tbl = countEachLabel(pxds)


%Visualize the pixel counts by class.

frequency = tbl.PixelCount/sum(tbl.PixelCount);

figure
bar(1:numel(classes),frequency)
xticks(1:numel(classes))
xticklabels(tbl.Name)
xtickangle(45)
ylabel('Frequency')

%%
%Resize CamVid Data
imageFolder = fullfile('imagesReszed',filesep);
imds = resizeCEImages(imds,imageFolder);

labelFolder = fullfile('labelsResized',filesep);
pxds = resizeCEPixelLabels(pxds,labelFolder);



%%
%Prepare Training and Test Sets
[imdsTrain, imdsTest, pxdsTrain, pxdsTest] = partitionCEData(imds,pxds);

numTrainingImages = numel(imdsTrain.Files)
numTestingImages = numel(imdsTest.Files)

%Create the network
imageSize = [256 256 3];
numClasses = numel(classes);
%lgraph = segnetLayers(imageSize,numClasses,'vgg16');


%%

%Balance Classes Using Class Weighting
% imageFreq = tbl.PixelCount ./ tbl.ImagePixelCount;
% classWeights = median(imageFreq) ./ imageFreq
% 
% pxLayer = pixelClassificationLayer('Name','labels','ClassNames', tbl.Name, 'ClassWeights', classWeights)
% 
% 
% lgraph = removeLayers(lgraph, 'pixelLabels');
% lgraph = addLayers(lgraph, pxLayer);
% lgraph = connectLayers(lgraph, 'softmax' ,'labels');

% load saved network architecture
load lgraph

%Select Training Options
options = trainingOptions('sgdm', ...
    'Momentum', 0.9, ...
    'InitialLearnRate', 1e-3, ...
    'L2Regularization', 0.0005, ...
    'MaxEpochs', 100, ...
    'MiniBatchSize', 3, ...
    'Shuffle', 'every-epoch', ...
    'VerboseFrequency', 2);



%%

%Data Augmentation
augmenter = imageDataAugmenter('RandXReflection',true,...
    'RandXTranslation', [-10 10], 'RandYTranslation',[-10 10]);


%Start Training
datasource = pixelLabelImageSource(imdsTrain,pxdsTrain, ...
                    'DataAugmentation',augmenter);

doTraining = false;
if doTraining
    [net, info] = trainNetwork(datasource,lgraph,options);
else
    data = load('CEtrainedSegNet.mat');
    net = data.net;
end

%%
%Test Network on One Image
tic
I = read(imdsTest);
C = semanticseg(I, net);

%Display the results.
B = labeloverlay(I, C, 'Colormap', cmap, 'Transparency',0.4);
figure
imshow(B)
pixelLabelColorbar(cmap, classes);

expectedResult = read(pxdsTest);
actual = uint8(C);
expected = uint8(expectedResult{1});
imshowpair(actual, expected)

iou = jaccard(C, expectedResult{1});
table(classes,iou)

%Evaluate Trained Network
pxdsResults = semanticseg(imdsTest,net,'WriteLocation',tempdir,'Verbose',false);
metrics = evaluateSemanticSegmentation(pxdsResults,pxdsTest,'Verbose',false);

toc