PfcSegmenter.hpp 2.9 KB
#pragma once

#include <array>
#include <functional>

#include <opencv2/imgproc.hpp>

#include "data/hdf5/CVMatLoader.hpp"

class PfcSegmenter {
  public:
	PfcSegmenter(const CVMatLoader &loader, const cv::Size &size) : size(size) {
		hostPfc = loader.asMat("/scene_model/PFC", CV_32F);
	}

	cv::Mat segment() const {
		return process(
		    CV_16UC1, [](const Mapping &mapping, cv::Mat &component) {
			    component.convertTo(component, CV_16UC1, (1L << 15) - 1);
			    return component & mapping.value;
		    });
	}

	void contour(std::vector<std::vector<std::vector<cv::Point>>> &contours,
	             cv::Mat &components) const {
		int i = 1;
		components =
		    process(CV_8UC1, [&i, &contours](const Mapping & /*mapping*/,
		                                     cv::Mat &component) {
			    cv::threshold(component, component, 0, i++, cv::THRESH_BINARY);

			    std::vector<std::vector<cv::Point>> contour;
			    cv::findContours(component, contour, cv::RETR_TREE,
			                     cv::CHAIN_APPROX_SIMPLE);

			    contours.emplace_back(std::move(contour));

			    return component;
		    });
	}

	cv::Mat mask(const std::string &label) const {
		return process(CV_8UC1, [&label](const Mapping &mapping,
		                                 cv::Mat &component) {
			if (label == mapping.label) {
				return component;
			}

			return cv::Mat{cv::Mat::zeros(component.size(), component.type())};
		});
	}

  private:
	const cv::Size size;
	cv::Mat hostPfc;

	class Mapping {
	  public:
		Mapping(int value, const cv::Range &range,
		        const std::string &label = "")
		    : value(value), range(range), label(label) {}

		const int value;
		const cv::Range range;
		const std::string label;
	};

	static const std::array<Mapping, 5> mappings;

	cv::Mat process(int outputType,
	                const std::function<cv::Mat(const Mapping &, cv::Mat &)>
	                    &modifyMask) const {
		cv::Mat output(size, outputType, cv::Scalar(0));

		const cv::Mat kernel =
		    cv::getStructuringElement(cv::MORPH_RECT, cv::Size(3, 1));

		cv::Mat component;
		for (const auto &mapping : mappings) {
			/* Workaround as the image is unsigned 32-bit thus it allows values
			 * above the signed limit. It is assumed that values greater than
			 * the signed maximum belongs to component, which end boundary is
			 * at the limit
			 */
			/* Threshold PFC identifier range */
			const auto upperLimit = 2'147'483'647;
			if (mapping.range.end >= upperLimit) {
				component = hostPfc >= upperLimit;
			} else {
				cv::inRange(hostPfc, mapping.range.start, mapping.range.end,
				            component);
			}
			/* Apply close morphology */
			cv::morphologyEx(component, component, cv::MORPH_CLOSE, kernel);
			/* Apply median filter */
			cv::medianBlur(component, component, 3);

			/* Set pixels to operational limit value */
			output = output | modifyMask(mapping, component);
		}

		/* Segmented PFC limits mask */
		return output;
	}
};