s = load("digitsCustom.mat");
dlnet = s.dlnet;
classes = s.classes;

digitDatasetPath = fullfile(matlabroot,'toolbox','nnet','nndemos', ...
    'nndatasets','DigitDataset');
imds = imageDatastore(digitDatasetPath, ...
    'IncludeSubfolders',true);

miniBatchSize = 128;
executionEnvironment = "auto";

imds.ReadSize = miniBatchSize;

numObservations = numel(imds.Files);
YPred = strings(1,numObservations);
i = 1;

% Loop over mini-batches.
while hasdata(imds)
    
    % Read mini-batch of data.
    data = read(imds);
    X = cat(4,data{:});
    
    % Normalize the images.
    X = single(X)/255;
    
    % Convert mini-batch of data to dlarray.
    dlX = dlarray(X,'SSCB');
    
    % If training on a GPU, then convert data to gpuArray.
    if (executionEnvironment == "auto" && canUseGPU) || executionEnvironment == "gpu"
        dlX = gpuArray(dlX);
    end
    
    % Make predictions using the predict function.
    dlYPred = predict(dlnet,dlX);
   
    % Determine corresponding classes.
    [~,idxTop] = max(extractdata(dlYPred),[],1);
    idxMiniBatch = i:min((i+miniBatchSize-1),numObservations);
    YPred(idxMiniBatch) = classes(idxTop);
    
    i = i + miniBatchSize;
end

idx = randperm(numObservations,9);
figure
for i = 1:9
    subplot(3,3,i)
    I = imread(imds.Files{idx(i)});
    label = YPred(idx(i));
    imshow(I)
    title("Label: " + label)
end