%% U-Net
% Skript pro načtení vstupních dat sítě U-Net, vytvoření validační množiny,
% vytvoření architektury neuronové sítě a její natrénování a testování.

clc; clear; close all

%% Načtení dat pro trénování/testování sítě
imgs_train_location = 'C:\Users\jana-\Documents\FEKT\DP\unet_final\unet_train\images';  % cesta ke složce obsahující trénovací obrazy
imgs_test_location = 'C:\Users\jana-\Documents\FEKT\DP\unet_final\unet_test\images';    % cesta ke složce obsahující testovací obrazy
lbls_train_location = 'C:\Users\jana-\Documents\FEKT\DP\unet_final\unet_train\labels';  % cesta ke složce obsahující trénovací obrazy s klasifikovanými cévami
lbls_test_location = 'C:\Users\jana-\Documents\FEKT\DP\unet_final\unet_test\labels';    % cesta ke složce obsahující testovací obrazy s klasifikovanými cévami
addpath(imgs_train_location);
addpath(imgs_test_location);
addpath(lbls_train_location);
addpath(lbls_test_location);

%% ImageDatastore
pom_imgs_train = imageDatastore(imgs_train_location);
imgs_test = imageDatastore(imgs_test_location);

% oddělení validační množiny od trénovací
pom_pocet_train = numel(pom_imgs_train.Files);         % počet souborů trénovací množiny
shuffle_index = randperm(pom_pocet_train);             % random zamíchání indexů
pocet_train = round(0.85 * pom_pocet_train);           % 85% dat bude trénovacích, zbytek dat bude použit pro validaci
train_index = shuffle_index(1:pocet_train);
valid_index = shuffle_index(pocet_train + 1:end);
pompom_imgs_train = pom_imgs_train.Files(train_index);
pompom_imgs_valid = pom_imgs_train.Files(valid_index);

imgs_train = imageDatastore(pompom_imgs_train);
imgs_valid = imageDatastore(pompom_imgs_valid);

%% PixelLabelDatastore
classNames = ["pozadi"
    "tepna"
    "zila"
    ];                   % definice klasifikačních tříd
pixelLabelID = [1 2 3];  % přiřazení numerické hodnoty jednotlivým třídám

pom_lbls_train = pixelLabelDatastore(lbls_train_location, classNames, pixelLabelID);
lbls_test = pixelLabelDatastore(lbls_test_location, classNames, pixelLabelID);

pompom_lbls_train = pom_lbls_train.Files(train_index);
pompom_lbls_valid = pom_lbls_train.Files(valid_index);

lbls_train = pixelLabelDatastore(pompom_lbls_train, classNames, pixelLabelID);
lbls_valid = pixelLabelDatastore(pompom_lbls_valid, classNames, pixelLabelID);

%% PixelLabelImageDatastore + augmentace trénovacích dat
augmenter = imageDataAugmenter(...
    'RandXReflection', true, ...
    'RandRotation',[-5, 5], ...
    'RandXScale',[0.9 1.1], ...
    'RandYScale',[0.9 1.1], ...
    'RandXTranslation',[-5 5], ...
    'RandYTranslation',[-5 5]);

pximds_train = pixelLabelImageDatastore(imgs_train, lbls_train, 'DataAugmentation', augmenter);
pximds_valid = pixelLabelImageDatastore(imgs_valid, lbls_valid);
pximds_test = pixelLabelImageDatastore(imgs_test, lbls_test);

% zobrazení vstupních dat (img a lbl) pomocí labeloverlay
% A = read(imgs_train);
% B = read(lbls_train);
% categories(B{1})
% C = labeloverlay(A, B{1});
% figure; imshow(C, []); title('Zobrazení labelů pixelů na snímku sítnice');

%% Výpočet vah jednotlivých tříd za účelem modifikace poslední vrstvy sítě
% přiřazení vah k jednotlivým třídám výpočtem invezní frekvence jednotlivých pixelů - váhovaná kategoriální cross entropie
tbl_pom = countEachLabel(pximds_train);
pocet_pix = sum(tbl_pom.PixelCount);
frekv = tbl_pom.PixelCount / pocet_pix;
frekv_inv = round((1./frekv) * 10) / 10;
frekv_inv = frekv_inv - 0.5;

%% Vytvoření architektury unet
lgraph = unetLayers([576 560 3], 3);   % vytvoření sítě U-Net, velikost vstupních obrazů = 576x560, počet tříd = 3
% plot(lgraph)                         % vykreslení architektury sítě
% analyzeNetwork(lgraph)               % analýza sítě
% save(lgraph)

% modifikace poslední vrstvy sítě na váhovanou křížovou entropii
new_pixelClass = pixelClassificationLayer('Name', 'new_Segmentation-Layer', 'Classes', tbl_pom.Name, 'ClassWeights', frekv_inv);
lgraph = replaceLayer(lgraph, 'Segmentation-Layer', new_pixelClass);

%% Trénování U-Net, nastavení trénovacích parametrů
options = trainingOptions('adam', ...
    'InitialLearnRate', 0.0001, ...
    'MaxEpochs', 30, ...
    'MiniBatchSize', 4, ...
    'Shuffle', 'every-epoch', ...
    'ValidationData', pximds_valid, ...
    'ValidationFrequency', 15, ...
    'Verbose', false, ...
    'ExecutionEnvironment', 'gpu', ...
    'Plots', 'training-progress');
trained_unet = trainNetwork(pximds_train, lgraph, options);

%% Testování naučené sítě U-Net
% load ('trained_unet.mat');      % načtení natrénované sítě

predikce = semanticseg(imgs_test, trained_unet, 'MiniBatchSize', 4, 'ExecutionEnvironment', 'gpu');
metriky = evaluateSemanticSegmentation(predikce, lbls_test);
