CVMatLoader.hpp 1.89 KB
#pragma once

#include <opencv2/core/mat.hpp>

#include "data/DataNotFoundException.hpp"
#include "data/hdf5/H5SingleDataReader.hpp"

class CVMatLoader {
  public:
	explicit CVMatLoader(const std::string &path) : path(path) {}

	cv::Mat asMat(const std::string &sourceDataSet, int destinationType) const {
		const auto reader =
		    createReader(sourceDataSet, resolveType(destinationType), 2);

		cv::Mat output(cv::Size{static_cast<int>(reader.dimensions[1]),
		                        static_cast<int>(reader.dimensions[0])},
		               destinationType);
		assert(output.isContinuous());

		reader.read(output.data);
		return output;
	}

	template <typename T>
	std::vector<T> asVec(const std::string &sourceDataSet,
	                     const H5::PredType &sourceType) const {
		const auto reader = createReader(sourceDataSet, sourceType, 1);

		std::vector<T> output(reader.size());
		reader.read(output.data());

		return output;
	}

  private:
	const std::string path;

	H5SingleDataReader createReader(const std::string &sourceDataSet,
	                                const H5::PredType &sourceType,
	                                int dimensions) const {
		try {
			return {path, sourceDataSet, sourceType, dimensions};
		} catch (const H5::FileIException &exception) {
			throw DataNotFoundException(exception.getDetailMsg());
		}
	}

	static const H5::PredType &resolveType(int cvType) {
		switch (CV_MAT_DEPTH(cvType)) {
		case CV_8U:
			return H5SingleDataReader::UINT8;
		case CV_16U:
			return H5SingleDataReader::UINT16;
		case CV_8S:
			return H5SingleDataReader::INT8;
		case CV_16S:
			return H5SingleDataReader::INT16;
		case CV_32S:
			return H5SingleDataReader::INT32;
		case CV_32F:
			return H5SingleDataReader::FLOAT;
		case CV_64F:
			return H5SingleDataReader::DOUBLE;
		default:
			throw std::runtime_error("Unknown OpenCV type " +
			                         std::to_string(cvType));
		}
	}
};