% 'cnn_digitTrain4D.m'
% Convolutional neural network for classification of the images in the
% built-in digitTrain4D dataset

% David Pan, UAH 

clear all;
close all;

% Load the training data, which contains 5000 images of digits. 
[XTrain,YTrain] = digitTrain4DArrayData;

% Show some sample images
figure;
perm = randperm(5000,64);
for i = 1:64
    subplot(8,8,i);
    imshow(XTrain(:,:,:,perm(i)));
end

% Set aside 1000 of the images for testing after the network is trained
idx = randperm(size(XTrain,4),1000);
XTest = XTrain(:,:,:,idx);
XTrain(:,:,:,idx) = [];
YTest = YTrain(idx);
YTrain(idx) = [];

% Set aside 1000 of the training images for validation during training.
idx = randperm(size(XTrain,4),1000);
XValidation = XTrain(:,:,:,idx);
XTrain(:,:,:,idx) = [];
YValidation = YTrain(idx);
YTrain(idx) = [];

layers = [
    imageInputLayer([28 28 1])
   
    convolution2dLayer(3,32)
    reluLayer
    maxPooling2dLayer(2, 'Stride',2)

    convolution2dLayer(3,64)
    reluLayer
    maxPooling2dLayer(2, 'Stride',2)
  
    fullyConnectedLayer(10)
    softmaxLayer
    classificationLayer];
 
options = trainingOptions('adam', ...
    'InitialLearnRate',0.01, ...
    'MaxEpochs',5, ...
    'Shuffle','every-epoch', ...
    'MiniBatchSize',64, ...
    'ValidationData', {XValidation,YValidation}, ...
    'ValidationFrequency', 20, ...
    'Verbose',false, ...
    'Plots','training-progress');

net = trainNetwork(XTrain,YTrain,layers,options);

YPred = classify(net,XTest);
accuracy = sum(YPred == YTest)/numel(YPred)

% Case study of a certain input image sample
im = XTrain(:,:,:,1);
figure; imshow(im)

act1 = activations(net,im,'conv_1');
I1 = imtile(act1, 'GridSize',[4 8]);
figure; imagesc(I1); colorbar

act2 = activations(net,im,'conv_2');
I2 = imtile(act2, 'GridSize',[8 8]);
figure; imagesc(I2); colorbar

% Output of the Softmax layer
act9 = activations(net,im,'softmax');
figure;
stem(reshape(act9(1,1,:), 10, 1));
grid

% Visualize the weights
W2 = net.Layers(2).Weights;
I_weight1 = imtile(W2, 'GridSize',[4 8]);
figure; imagesc(I_weight1); colorbar