summaryrefslogtreecommitdiff
path: root/tests/MnistDatabase.cpp
blob: 5c10b0c2b4ac13ac256f22babb13bffac29f8319 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
//
// Copyright © 2017 Arm Ltd. All rights reserved.
// See LICENSE file in the project root for full license information.
//
#include "MnistDatabase.hpp"

#include <boost/numeric/conversion/cast.hpp>
#include <boost/log/trivial.hpp>
#include <boost/assert.hpp>
#include <fstream>
#include <vector>

constexpr int g_kMnistImageByteSize = 28 * 28;

void EndianSwap(unsigned int &x)
{
    x = (x >> 24) | ((x << 8) & 0x00FF0000) | ((x >> 8) & 0x0000FF00) | (x << 24);
}

MnistDatabase::MnistDatabase(const std::string& binaryFileDirectory, bool scaleValues)
    : m_BinaryDirectory(binaryFileDirectory)
    , m_ScaleValues(scaleValues)
{
}

std::unique_ptr<MnistDatabase::TTestCaseData> MnistDatabase::GetTestCaseData(unsigned int testCaseId)
{
    std::vector<unsigned char> I(g_kMnistImageByteSize);
    unsigned int label = 0;

    std::string imagePath = m_BinaryDirectory + std::string("t10k-images.idx3-ubyte");
    std::string labelPath = m_BinaryDirectory + std::string("t10k-labels.idx1-ubyte");

    std::ifstream imageStream(imagePath, std::ios::binary);
    std::ifstream labelStream(labelPath, std::ios::binary);

    if (!imageStream.is_open())
    {
        BOOST_LOG_TRIVIAL(fatal) << "Failed to load " << imagePath;
        return nullptr;
    }
    if (!labelStream.is_open())
    {
        BOOST_LOG_TRIVIAL(fatal) << "Failed to load " << imagePath;
        return nullptr;
    }

    unsigned int magic, num, row, col;

    // check the files have the correct header
    imageStream.read(reinterpret_cast<char*>(&magic), sizeof(magic));
    if (magic != 0x03080000)
    {
        BOOST_LOG_TRIVIAL(fatal) << "Failed to read " << imagePath;
        return nullptr;
    }
    labelStream.read(reinterpret_cast<char*>(&magic), sizeof(magic));
    if (magic != 0x01080000)
    {
        BOOST_LOG_TRIVIAL(fatal) << "Failed to read " << labelPath;
        return nullptr;
    }

    // Endian swap image and label file - All the integers in the files are stored in MSB first(high endian) format,
    // hence need to flip the bytes of the header if using it on Intel processors or low-endian machines
    labelStream.read(reinterpret_cast<char*>(&num), sizeof(num));
    imageStream.read(reinterpret_cast<char*>(&num), sizeof(num));
    EndianSwap(num);
    imageStream.read(reinterpret_cast<char*>(&row), sizeof(row));
    EndianSwap(row);
    imageStream.read(reinterpret_cast<char*>(&col), sizeof(col));
    EndianSwap(col);

    // read image and label into memory
    imageStream.seekg(testCaseId * g_kMnistImageByteSize, std::ios_base::cur);
    imageStream.read(reinterpret_cast<char*>(&I[0]), g_kMnistImageByteSize);
    labelStream.seekg(testCaseId, std::ios_base::cur);
    labelStream.read(reinterpret_cast<char*>(&label), 1);

    if (!imageStream.good())
    {
        BOOST_LOG_TRIVIAL(fatal) << "Failed to read " << imagePath;
        return nullptr;
    }
    if (!labelStream.good())
    {
        BOOST_LOG_TRIVIAL(fatal) << "Failed to read " << labelPath;
        return nullptr;
    }

    std::vector<float> inputImageData;
    inputImageData.resize(g_kMnistImageByteSize);

    for (unsigned int i = 0; i < col * row; ++i)
    {
        inputImageData[i] = boost::numeric_cast<float>(I[i]);

        if(m_ScaleValues)
        {
            inputImageData[i] /= 255.0f;
        }
    }

    return std::make_unique<TTestCaseData>(label, std::move(inputImageData));
}