Fix and clean up ImageTensorGenerator
* Fix the issue that uint8_t image tensors are written out as
characters.
* Add doxygen comments.
Signed-off-by: SiCong Li <sicong.li@arm.com>
Change-Id: Ic5b3af489f405abcc1ed1e7dd76f56c5ae7ffbf4
diff --git a/tests/ImageTensorGenerator/ImageTensorGenerator.hpp b/tests/ImageTensorGenerator/ImageTensorGenerator.hpp
index e5bc414..3e164bc 100644
--- a/tests/ImageTensorGenerator/ImageTensorGenerator.hpp
+++ b/tests/ImageTensorGenerator/ImageTensorGenerator.hpp
@@ -12,6 +12,7 @@
#include <iterator>
#include <string>
+// Parameters used in normalizing images
struct NormalizationParameters
{
float scale{ 1.0 };
@@ -26,9 +27,13 @@
TFLite = 2,
};
-// Get normalization parameters.
-// Note that different flavours of models have different normalization methods.
-// This tool currently only supports Caffe, TF and TFLite models
+/** Get normalization parameters.
+ * Note that different flavours of models and different model data types have different normalization methods.
+ * This tool currently only supports Caffe, TF and TFLite models
+ *
+ * @param[in] modelFormat One of the supported frontends
+ * @param[in] outputType Output type of the image tensor, also the type of the intended model
+ */
NormalizationParameters GetNormalizationParameters(const SupportedFrontend& modelFormat,
const armnn::DataType& outputType)
{
@@ -62,7 +67,15 @@
return normParams;
}
-// Prepare raw image tensor data by loading the image from imagePath and preprocessing it.
+/** Prepare raw image tensor data by loading the image from imagePath and preprocessing it.
+ *
+ * @param[in] imagePath Path to the image file
+ * @param[in] newWidth The new width of the output image tensor
+ * @param[in] newHeight The new height of the output image tensor
+ * @param[in] normParams Normalization parameters for the normalization of the image
+ * @param[in] batchSize Batch size
+ * @param[in] outputLayout Data layout of the output image tensor
+ */
template <typename ElemType>
std::vector<ElemType> PrepareImageTensor(const std::string& imagePath,
unsigned int newWidth,
@@ -148,9 +161,21 @@
return imageDataQasymm8;
}
-// Write image tensor to ofstream
+/** Write image tensor to ofstream
+ *
+ * @param[in] imageData Image tensor data
+ * @param[in] imageTensorFile Output filestream (ofstream) to which the image tensor data is written
+ */
template <typename ElemType>
void WriteImageTensorImpl(const std::vector<ElemType>& imageData, std::ofstream& imageTensorFile)
{
std::copy(imageData.begin(), imageData.end(), std::ostream_iterator<ElemType>(imageTensorFile, " "));
+}
+
+// For uint8_t image tensor, cast it to int before writing it to prevent writing data as characters instead of
+// numerical values
+template <>
+void WriteImageTensorImpl<uint8_t>(const std::vector<uint8_t>& imageData, std::ofstream& imageTensorFile)
+{
+ std::copy(imageData.begin(), imageData.end(), std::ostream_iterator<int>(imageTensorFile, " "));
}
\ No newline at end of file