tests/MnistDatabase.cpp

   1 //
   2 // Copyright © 2017 Arm Ltd. All rights reserved.
   3 // SPDX-License-Identifier: MIT
   4 //
   5 #include "MnistDatabase.hpp"
   6
   7 #include <armnn/Logging.hpp>
   8
   9 #include <boost/numeric/conversion/cast.hpp>
  10 #include <boost/assert.hpp>
  11 #include <fstream>
  12 #include <vector>
  13
  14 constexpr int g_kMnistImageByteSize = 28 * 28;
  15
  16 void EndianSwap(unsigned int &x)
  17 {
  18     x = (x >> 24) | ((x << 8) & 0x00FF0000) | ((x >> 8) & 0x0000FF00) | (x << 24);
  19 }
  20
  21 MnistDatabase::MnistDatabase(const std::string& binaryFileDirectory, bool scaleValues)
  22     : m_BinaryDirectory(binaryFileDirectory)
  23     , m_ScaleValues(scaleValues)
  24 {
  25 }
  26
  27 std::unique_ptr<MnistDatabase::TTestCaseData> MnistDatabase::GetTestCaseData(unsigned int testCaseId)
  28 {
  29     std::vector<unsigned char> I(g_kMnistImageByteSize);
  30     unsigned int label = 0;
  31
  32     std::string imagePath = m_BinaryDirectory + std::string("t10k-images.idx3-ubyte");
  33     std::string labelPath = m_BinaryDirectory + std::string("t10k-labels.idx1-ubyte");
  34
  35     std::ifstream imageStream(imagePath, std::ios::binary);
  36     std::ifstream labelStream(labelPath, std::ios::binary);
  37
  38     if (!imageStream.is_open())
  39     {
  40         ARMNN_LOG(fatal) << "Failed to load " << imagePath;
  41         return nullptr;
  42     }
  43     if (!labelStream.is_open())
  44     {
  45         ARMNN_LOG(fatal) << "Failed to load " << imagePath;
  46         return nullptr;
  47     }
  48
  49     unsigned int magic, num, row, col;
  50
  51     // Checks the files have the correct header.
  52     imageStream.read(reinterpret_cast<char*>(&magic), sizeof(magic));
  53     if (magic != 0x03080000)
  54     {
  55         ARMNN_LOG(fatal) << "Failed to read " << imagePath;
  56         return nullptr;
  57     }
  58     labelStream.read(reinterpret_cast<char*>(&magic), sizeof(magic));
  59     if (magic != 0x01080000)
  60     {
  61         ARMNN_LOG(fatal) << "Failed to read " << labelPath;
  62         return nullptr;
  63     }
  64
  65     // Endian swaps the image and label file - all the integers in the files are stored in MSB first(high endian)
  66     // format, hence it needs to flip the bytes of the header if using it on Intel processors or low-endian machines
  67     labelStream.read(reinterpret_cast<char*>(&num), sizeof(num));
  68     imageStream.read(reinterpret_cast<char*>(&num), sizeof(num));
  69     EndianSwap(num);
  70     imageStream.read(reinterpret_cast<char*>(&row), sizeof(row));
  71     EndianSwap(row);
  72     imageStream.read(reinterpret_cast<char*>(&col), sizeof(col));
  73     EndianSwap(col);
  74
  75     // Reads image and label into memory.
  76     imageStream.seekg(testCaseId * g_kMnistImageByteSize, std::ios_base::cur);
  77     imageStream.read(reinterpret_cast<char*>(&I[0]), g_kMnistImageByteSize);
  78     labelStream.seekg(testCaseId, std::ios_base::cur);
  79     labelStream.read(reinterpret_cast<char*>(&label), 1);
  80
  81     if (!imageStream.good())
  82     {
  83         ARMNN_LOG(fatal) << "Failed to read " << imagePath;
  84         return nullptr;
  85     }
  86     if (!labelStream.good())
  87     {
  88         ARMNN_LOG(fatal) << "Failed to read " << labelPath;
  89         return nullptr;
  90     }
  91
  92     std::vector<float> inputImageData;
  93     inputImageData.resize(g_kMnistImageByteSize);
  94
  95     for (unsigned int i = 0; i < col * row; ++i)
  96     {
  97         inputImageData[i] = boost::numeric_cast<float>(I[i]);
  98
  99         if(m_ScaleValues)
 100         {
 101             inputImageData[i] /= 255.0f;
 102         }
 103     }
 104
 105     return std::make_unique<TTestCaseData>(label, std::move(inputImageData));
 106 }