Train an OCR Model to Recognize Seven-Segment Digits

datasetFiles = helperDownloadDataset;

ld = load("7SegmentGtruth.mat");
gTruth = ld.gTruth;

labelName = "Text";
attributeName = "Digits";
[imds,boxds,txtds] = ocrTrainingData(gTruth,labelName,attributeName);

helperDisplayGroundtruthData(imds, boxds, txtds)

allImagesText = txtds.readall;
allText = strjoin(vertcat(allImagesText{:}), "");

[characterSet, ~, idx] = unique(char(allText));

disp("Ground truth Character Set: " + string(characterSet))

characterSet = cellstr(characterSet');
characterCount = accumarray(idx,1);

characterCountTbl = table(characterSet, characterCount, ...
    VariableNames=["Character", "Count"]);
characterCountTbl = sortrows(characterCountTbl, ...
    "Count", "descend")

numCharacters = numel(characterSet);

figure
bar(1:numCharacters, characterCountTbl.Count)
xticks(1:numCharacters)
xticklabels(characterCountTbl.Character)
xlabel("Digits")
ylabel("Number of samples")

cds = combine(imds,boxds,txtds);

trainPercent = 60;
[cdsTrain, cdsVal, cdsTest, numTrain, numVal, numTest] = helperPartitionOCRData(cds, trainPercent);

disp("Number of training images = " + numTrain)

disp("Number of validation images = " + numVal)

disp("Number of test images = " + numTest)

outputDir = "OCRModel";
if ~exist(outputDir, "dir")
    mkdir(outputDir);
end

checkpointsDir = "Checkpoints";
if ~exist(checkpointsDir, "dir")
    mkdir(checkpointsDir);
end

ocrOptions = ocrTrainingOptions(GradientDecayFactor=0.9,...
    InitialLearnRate=20e-4,...
    MaxEpochs=15,...
    VerboseFrequency=100,...
    OutputLocation=outputDir,...
    CheckpointPath=checkpointsDir,...
    ValidationData=cdsVal,...
    ValidationFrequency=10);

trainedModelName = "sevenSegmentModel";
baseModel = "english";
[trainedModel, trainingInfo] = trainOCR(cdsTrain, trainedModelName, baseModel, ocrOptions);

figure
plot(trainingInfo.TrainingRMSE); hold on;
plot(trainingInfo.ValidationRMSE)
legend(["Training", "Validation"])
xlabel("Iterations")
ylabel("RMSE")
title("Training vs Validation RMSE Curve")

trainedModelResults = ocr(cdsTest,Model=trainedModel);
trainedModelMetrics = evaluateOCR(trainedModelResults,cdsTest);

trainedModelAccuracy = 100*(1-trainedModelMetrics.DataSetMetrics.CharacterErrorRate);
disp("Test accuracy of the trained model= " + trainedModelAccuracy + "%")

I = imread("sevSegDisp.jpg");
roi = [506 725 1418 626];
ocrResults = ocr(I,roi,Model=trainedModel,LayoutAnalysis="Block");
Iocr = insertObjectAnnotation(I,"rectangle",...
            ocrResults.WordBoundingBoxes, ...
            ocrResults.Words, LineWidth=5,FontSize=72);
figure
imshow(Iocr)

quantizedModelName = "quantizedModel";
quantizedModel = quantizeOCR(trainedModel,quantizedModelName);

fOCR = @() ocr(I, Model=trainedModel);
tOCR = timeit(fOCR);
 
fQuantizedOCR = @() ocr(I, Model=quantizedModel);
tQuantizedOCR = timeit(fQuantizedOCR);

perfRatio = tOCR/tQuantizedOCR;
disp("Quantized model is " + perfRatio + "x faster");

trainedModelFile         = dir(trainedModel);
trainedModelFileSizeInMB = trainedModelFile.bytes/1000000;

quantizedModelFile         = dir(quantizedModel);
quantizedModelFileSizeInMB = quantizedModelFile.bytes/1000000;

sizeRatio = trainedModelFileSizeInMB/quantizedModelFileSizeInMB;
disp("Quantized model is " + sizeRatio + "x smaller");

quantizedModelResults = ocr(cdsTest,Model=quantizedModel);
quantizedModelMetrics = evaluateOCR(quantizedModelResults,cdsTest);

quantizedModelAccuracy = 100*(1-quantizedModelMetrics.DataSetMetrics.CharacterErrorRate);
disp("Test accuracy of the quantized model = " + quantizedModelAccuracy + "%")

dropInAccuracy = trainedModelAccuracy - quantizedModelAccuracy;
disp("Drop in accuracy after quantization = " + dropInAccuracy + "%")

trainedModelResults = [trainedModelAccuracy; trainedModelFileSizeInMB; tOCR];
quantizedModelResults = [quantizedModelAccuracy; quantizedModelFileSizeInMB; tQuantizedOCR];

table(trainedModelResults, quantizedModelResults, ...
    VariableNames=[trainedModelName, quantizedModelName], ...
    RowNames=["Accuracy (in %)", "File Size (in MB)", "Runtime (in seconds)"])

function datasetFiles = helperDownloadDataset()

    datasetURL = "https://ssd.mathworks.com/supportfiles/vision/data/7SegmentImages.zip";
    datasetZip = "7SegmentImages.zip";
    if ~exist(datasetZip,"file")
        disp("Downloading evaluation data set (" + datasetZip + " - 96 MB)...");
        websave(datasetZip,datasetURL);
    end
    
    datasetFiles = unzip(datasetZip);
end

function helperDisplayGroundtruthData(imds, boxds, txtds)

    figure("Position", [10 10 900 600])
    tiledlayout(2,2,TileSpacing="tight",Padding="tight")
    
    for i = 1:4
        nexttile
        img = read(imds);
        bbox = read(boxds);
        label = read(txtds);
        img = insertObjectAnnotation(img,"rectangle",...
            bbox{1}, label{1}, LineWidth=15, FontSize=72, TextBoxOpacity=0.9);
        imshow(img);
    end

    reset(imds);
    reset(boxds);
    reset(txtds);
end

function [cdsTrain, cdsVal, cdsTest, numTrain, numVal, numTest] = helperPartitionOCRData(cds, trainPercent)

    % Set initial random state for example reproducibility.
    rng(0);
    
    % Shuffle the sample order in the dataset.
    imds = cds.UnderlyingDatastores{1};
    numSamples = numel(imds.Files);
    shuffledIndices = randperm(numSamples);
    
    % Use trainPercent of samples for training.
    trainRatio = trainPercent/100;
    numTrain = round(trainRatio*numSamples);
    trainIndices = shuffledIndices(1:numTrain);
    cdsTrain = subset(cds, trainIndices);
    
    % Split the rest of the samples evenly for validation and testing.
    numRest = numSamples - numTrain;
    numVal = ceil(numRest/2);
    numTest = numRest - numVal;
    
    valIndices = shuffledIndices(numTrain+1:numTrain+numVal);
    testIndices = shuffledIndices(numTrain+numVal+1:end);
    
    cdsVal = subset(cds, valIndices);
    cdsTest = subset(cds, testIndices);
end