clear all;

%digitDatasetPath = {...
%    '../data/16x16_bw/rasen-ball-robot-move-5-best-24x24',...
%    '../data/16x16_bw/2016-06-29-patches-play' };

digitDatasetPath = {...
    '../data/data-set-10'};

digitData = imageDatastore(digitDatasetPath, ...
        'IncludeSubfolders',true,...
        'LabelSource','foldernames', 'ReadFcn',@imgReader);


%%
% The labels for each patch is automatically determined by the name of the
% subfolder from digitDatasetPath the image is located at
countLabel = digitData.countEachLabel;
disp(countLabel)
numBallPatches = countLabel.Count(1)
%%
trainingNumFiles = round(numBallPatches*0.75);
%rng(1) % For reproducibility
[trainDigitData,testDigitData] = splitEachLabel(digitData, ...
				trainingNumFiles,'randomize');

layers = [imageInputLayer([16 16 1], 'DataAugmentation',{'randfliplr', 'randcrop'}) % ,'randcrop'
          %dropoutLayer()
          convolution2dLayer(5, 4, 'Stride', 2, 'Padding', 2)
          reluLayer
          maxPooling2dLayer(2, 'Stride', 2)
          %dropoutLayer()
          convolution2dLayer(3, 10, 'Stride', 1, 'Padding', 1)
          reluLayer
          maxPooling2dLayer(2, 'Stride', 1)
          %dropoutLayer()
          fullyConnectedLayer(32)
          reluLayer
          fullyConnectedLayer(2)
          softmaxLayer
          classificationLayer];

% Initialize the first convolutional layer weights using normally distributed 
% random numbers with standard deviation of 0.0001.
layer = layers(2);
N = (16*16 + 16*16*4)*0.5;
layer.Weights = 1/N * randn(layer.FilterSize(1), layer.FilterSize(2),1,layer.NumFilters);

layer = layers(5);
N = (8*8 + 8*8*10)*0.5;
layer.Weights = 1/N * randn(layer.FilterSize(1), layer.FilterSize(2),layers(2).NumFilters,layer.NumFilters);


%https://www.mathworks.com/help/vision/examples/object-detection-using-deep-learning.html
%{
opts = trainingOptions('sgdm', ...
    'Momentum', 0.9, ...
    'InitialLearnRate', 0.001, ...
    'LearnRateSchedule', 'piecewise', ...
    'LearnRateDropFactor', 0.1, ...
    'LearnRateDropPeriod', 8, ...
    'L2Regularization', 0.004, ...
    'MaxEpochs', 40, ...
    'MiniBatchSize', 128, ...
    'Verbose', true);
%}

options = trainingOptions('sgdm',...
    'Momentum', 0.9, ...
    'InitialLearnRate', 0.001, ...
    'LearnRateSchedule', 'piecewise', ...
    'LearnRateDropFactor', 0.1, ...
    'LearnRateDropPeriod', 8, ...
    'L2Regularization', 0.004, ...
    'MaxEpochs', 30, ...
    'MiniBatchSize', 128); % 'cpu', 'parallel'

convnet = trainNetwork(digitData,layers,options);

save cn17.mat convnet

