ImageTensorGenerator.cpp File Reference

#include "ImageTensorGenerator.hpp"
#include "../InferenceTestImage.hpp"
#include <armnn/Logging.hpp>
#include <armnn/TypesUtils.hpp>
#include <boost/filesystem.hpp>
#include <boost/filesystem/operations.hpp>
#include <boost/filesystem/path.hpp>
#include <boost/program_options.hpp>
#include <boost/variant.hpp>
#include <algorithm>
#include <fstream>
#include <iostream>
#include <string>

Go to the source code of this file.

Functions
int	main (int argc, char *argv[])

Function Documentation

main()

int main	(	int	argc,
		char *	argv[]
	)

Definition at line 248 of file ImageTensorGenerator.cpp.

References ARMNN_LOG, armnn::Float32, GetNormalizationParameters(), PrepareImageTensor< float >(), PrepareImageTensor< int >(), PrepareImageTensor< uint8_t >(), armnn::QAsymmU8, armnn::Signed32, Exception::what(), and WriteImageTensorImpl().

{
    CommandLineProcessor cmdline;
    if (!cmdline.ProcessCommandLine(argc, argv))
    {
        return -1;
    }
    const std::string imagePath(cmdline.GetInputFileName());
    const std::string outputPath(cmdline.GetOutputFileName());
    const SupportedFrontend& modelFormat(cmdline.GetModelFormat());
    const armnn::DataType outputType(cmdline.GetOutputType());
    const unsigned int newWidth  = cmdline.GetNewWidth();
    const unsigned int newHeight = cmdline.GetNewHeight();
    const unsigned int batchSize = 1;
    const armnn::DataLayout outputLayout(cmdline.GetLayout());

    using TContainer = boost::variant<std::vector<float>, std::vector<int>, std::vector<uint8_t>>;
    std::vector<TContainer> imageDataContainers;
    const NormalizationParameters& normParams = GetNormalizationParameters(modelFormat, outputType);
    try
    {
        switch (outputType)
        {
            case armnn::DataType::Signed32:
                imageDataContainers.push_back(PrepareImageTensor<int>(
                    imagePath, newWidth, newHeight, normParams, batchSize, outputLayout));
                break;
            case armnn::DataType::QAsymmU8:
                imageDataContainers.push_back(PrepareImageTensor<uint8_t>(
                    imagePath, newWidth, newHeight, normParams, batchSize, outputLayout));
                break;
            case armnn::DataType::Float32:
            default:
                imageDataContainers.push_back(PrepareImageTensor<float>(
                    imagePath, newWidth, newHeight, normParams, batchSize, outputLayout));
                break;
        }
    }
    catch (const InferenceTestImageException& e)
    {
        ARMNN_LOG(fatal) << "Failed to load image file " << imagePath << " with error: " << e.what();
        return -1;
    }

    std::ofstream imageTensorFile;
    imageTensorFile.open(outputPath, std::ofstream::out);
    if (imageTensorFile.is_open())
    {
        boost::apply_visitor([&imageTensorFile](auto&& imageData) { WriteImageTensorImpl(imageData, imageTensorFile); },
                             imageDataContainers[0]);
        if (!imageTensorFile)
        {
            ARMNN_LOG(fatal) << "Failed to write to output file" << outputPath;
            imageTensorFile.close();
            return -1;
        }
        imageTensorFile.close();
    }
    else
    {
        ARMNN_LOG(fatal) << "Failed to open output file" << outputPath;
        return -1;
    }

    return 0;
}