#include "superScale.h"


#define CLIP(x, x1, x2) cv::max(x1, cv::min(x, x2))

int SuperScale::init(const std::string &proto_path, const std::string &model_path) {
	srnet_ = cv::dnn::readNetFromCaffe(proto_path, model_path);
	net_loaded_ = true;
	return 0;
}

cv::Mat SuperScale::processImageScale(const cv::Mat &src, float scale, const bool &use_sr,
	int sr_max_size) {
	cv::Mat dst = src;
	if (scale == 1.0) {  // src
		return dst;
	}

	int width = src.cols;
	int height = src.rows;
	if (scale == 2.0) {  // upsample
		int SR_TH = sr_max_size;
		if (use_sr && (int)sqrt(width * height * 1.0) < SR_TH && net_loaded_) {
			int ret = superResoutionScale(src, dst);
			if (ret == 0) return dst;
		}

		{ resize(src, dst, cv::Size(), scale, scale, cv::INTER_CUBIC); }
	}
	else if (scale < 1.0) {  // downsample
		resize(src, dst, cv::Size(), scale, scale, cv::INTER_AREA);
	}

	return dst;
}

int SuperScale::superResoutionScale(const cv::Mat &src, cv::Mat &dst) {
	cv::Mat blob;
	cv::dnn::blobFromImage(src, blob, 1.0 / 255, cv::Size(src.cols, src.rows), { 0.0f }, false, false);

	srnet_.setInput(blob);
	auto prob = srnet_.forward();

	dst = cv::Mat(prob.size[2], prob.size[3], CV_8UC1);

	for (int row = 0; row < prob.size[2]; row++) {
		const float *prob_score = prob.ptr<float>(0, 0, row);
		for (int col = 0; col < prob.size[3]; col++) {
			float pixel = prob_score[col] * 255.0f;
			dst.at<uint8_t>(row, col) = static_cast<uint8_t>(CLIP(pixel, 0.0f, 255.0f));
		}
	}
	return 0;
}