1 files changed, 499 insertions, 0 deletions
diff --git a/src/ipa/libipa/awb_bayes.cpp b/src/ipa/libipa/awb_bayes.cpp
new file mode 100644
index 00000000..e75bfcd6
--- /dev/null
+++ b/src/ipa/libipa/awb_bayes.cpp
@@ -0,0 +1,499 @@
+/* SPDX-License-Identifier: LGPL-2.1-or-later */
+/*
+ * Copyright (C) 2019, Raspberry Pi Ltd
+ * Copyright (C) 2024 Ideas on Board Oy
+ *
+ * Implementation of a bayesian AWB algorithm
+ */
+
+#include "awb_bayes.h"
+
+#include <cmath>
+
+#include <libcamera/base/log.h>
+#include <libcamera/control_ids.h>
+
+#include "colours.h"
+
+/**
+ * \file awb_bayes.h
+ * \brief Implementation of bayesian auto white balance algorithm
+ *
+ * This implementation is based on the initial implementation done by
+ * RaspberryPi.
+ * \todo: Documentation
+ *
+ * \todo Not all the features implemented by RaspberryPi were ported over to
+ * this algorithm because they either rely on hardware features not generally
+ * available or were considered not important enough at the moment.
+ *
+ * The following parts are not implemented:
+ *
+ * - min_pixels: minimum proportion of pixels counted within AWB region for it
+ *   to be "useful"
+ * - min_g: minimum G value of those pixels, to be regarded a "useful"
+ * - min_regions: number of AWB regions that must be "useful" in order to do the
+ *   AWB calculation
+ * - deltaLimit: clamp on colour error term (so as not to penalize non-grey
+ *   excessively)
+ * - bias_proportion: The biasProportion parameter adds a small proportion of
+ *   the counted pixels to a region biased to the biasCT colour temperature.
+ *   A typical value for biasProportion would be between 0.05 to 0.1.
+ * - bias_ct: CT target for the search bias
+ * - sensitivityR: red sensitivity ratio (set to canonical sensor's R/G divided
+ *   by this sensor's R/G)
+ * - sensitivityB: blue sensitivity ratio (set to canonical sensor's B/G divided
+ *   by this sensor's B/G)
+ */
+
+namespace libcamera {
+
+LOG_DECLARE_CATEGORY(Awb)
+
+namespace {
+
+template<typename T>
+class LimitsRecorder
+{
+public:
+	LimitsRecorder()
+		: min_(std::numeric_limits<T>::max()),
+		  max_(std::numeric_limits<T>::min())
+	{
+	}
+
+	void record(const T &value)
+	{
+		min_ = std::min(min_, value);
+		max_ = std::max(max_, value);
+	}
+
+	const T &min() const { return min_; }
+	const T &max() const { return max_; }
+
+private:
+	T min_;
+	T max_;
+};
+
+#ifndef __DOXYGEN__
+template<typename T>
+std::ostream &operator<<(std::ostream &out, const LimitsRecorder<T> &v)
+{
+	out << "[ " << v.min() << ", " << v.max() << " ]";
+	return out;
+}
+#endif
+
+} /* namespace */
+
+namespace ipa {
+
+/**
+ * \brief Step size control for CT search
+ */
+constexpr double kSearchStep = 0.2;
+
+/**
+ * \copydoc libcamera::ipa::Interpolator::interpolate()
+ */
+template<>
+void Interpolator<Pwl>::interpolate(const Pwl &a, const Pwl &b, Pwl &dest, double lambda)
+{
+	dest = Pwl::combine(a, b,
+			    [=](double /*x*/, double y0, double y1) -> double {
+				    return y0 * (1.0 - lambda) + y1 * lambda;
+			    });
+}
+
+/**
+ * \class AwbBayes
+ * \brief Implementation of a bayesian auto white balance algorithm
+ *
+ * In a bayesian AWB algorithm the auto white balance estimation is improved by
+ * taking the likelihood of a given lightsource based on the estimated lux level
+ * into account. E.g. If it is very bright we can assume that we are outside and
+ * that colour temperatures around 6500 are preferred.
+ *
+ * The second part of this algorithm is the search for the most likely colour
+ * temperature. It is implemented in AwbBayes::coarseSearch() and in
+ * AwbBayes::fineSearch(). The search works very well without prior likelihoods
+ * and therefore the algorithm itself provides very good results even without
+ * prior likelihoods.
+ */
+
+/**
+ * \var AwbBayes::transversePos_
+ * \brief How far to wander off CT curve towards "more purple"
+ */
+
+/**
+ * \var AwbBayes::transverseNeg_
+ * \brief How far to wander off CT curve towards "more green"
+ */
+
+/**
+ * \var AwbBayes::currentMode_
+ * \brief The currently selected mode
+ */
+
+int AwbBayes::init(const YamlObject &tuningData)
+{
+	int ret = colourGainCurve_.readYaml(tuningData["colourGains"], "ct", "gains");
+	if (ret) {
+		LOG(Awb, Error)
+			<< "Failed to parse 'colourGains' "
+			<< "parameter from tuning file";
+		return ret;
+	}
+
+	ctR_.clear();
+	ctB_.clear();
+	for (const auto &[ct, g] : colourGainCurve_.data()) {
+		ctR_.append(ct, 1.0 / g[0]);
+		ctB_.append(ct, 1.0 / g[1]);
+	}
+
+	/* We will want the inverse functions of these too. */
+	ctRInverse_ = ctR_.inverse().first;
+	ctBInverse_ = ctB_.inverse().first;
+
+	ret = readPriors(tuningData);
+	if (ret) {
+		LOG(Awb, Error) << "Failed to read priors";
+		return ret;
+	}
+
+	ret = parseModeConfigs(tuningData, controls::AwbAuto);
+	if (ret) {
+		LOG(Awb, Error)
+			<< "Failed to parse mode parameter from tuning file";
+		return ret;
+	}
+	currentMode_ = &modes_[controls::AwbAuto];
+
+	transversePos_ = tuningData["transversePos"].get<double>(0.01);
+	transverseNeg_ = tuningData["transverseNeg"].get<double>(0.01);
+	if (transversePos_ <= 0 || transverseNeg_ <= 0) {
+		LOG(Awb, Error) << "AwbConfig: transversePos/Neg must be > 0";
+		return -EINVAL;
+	}
+
+	return 0;
+}
+
+int AwbBayes::readPriors(const YamlObject &tuningData)
+{
+	const auto &priorsList = tuningData["priors"];
+	std::map<uint32_t, Pwl> priors;
+
+	if (!priorsList) {
+		LOG(Awb, Error) << "No priors specified";
+		return -EINVAL;
+	}
+
+	for (const auto &p : priorsList.asList()) {
+		if (!p.contains("lux")) {
+			LOG(Awb, Error) << "Missing lux value";
+			return -EINVAL;
+		}
+
+		uint32_t lux = p["lux"].get<uint32_t>(0);
+		if (priors.count(lux)) {
+			LOG(Awb, Error) << "Duplicate prior for lux value " << lux;
+			return -EINVAL;
+		}
+
+		std::vector<uint32_t> temperatures =
+			p["ct"].getList<uint32_t>().value_or(std::vector<uint32_t>{});
+		std::vector<double> probabilities =
+			p["probability"].getList<double>().value_or(std::vector<double>{});
+
+		if (temperatures.size() != probabilities.size()) {
+			LOG(Awb, Error)
+				<< "Ct and probability array sizes are unequal";
+			return -EINVAL;
+		}
+
+		if (temperatures.empty()) {
+			LOG(Awb, Error)
+				<< "Ct and probability arrays are empty";
+			return -EINVAL;
+		}
+
+		std::map<int, double> ctToProbability;
+		for (unsigned int i = 0; i < temperatures.size(); i++) {
+			int t = temperatures[i];
+			if (ctToProbability.count(t)) {
+				LOG(Awb, Error) << "Duplicate ct value";
+				return -EINVAL;
+			}
+
+			ctToProbability[t] = probabilities[i];
+		}
+
+		auto &pwl = priors[lux];
+		for (const auto &[ct, prob] : ctToProbability) {
+			if (prob < 1e-6) {
+				LOG(Awb, Error) << "Prior probability must be larger than 1e-6";
+				return -EINVAL;
+			}
+			pwl.append(ct, prob);
+		}
+	}
+
+	if (priors.empty()) {
+		LOG(Awb, Error) << "No priors specified";
+		return -EINVAL;
+	}
+
+	priors_.setData(std::move(priors));
+
+	return 0;
+}
+
+void AwbBayes::handleControls(const ControlList &controls)
+{
+	auto mode = controls.get(controls::AwbMode);
+	if (mode) {
+		auto it = modes_.find(static_cast<controls::AwbModeEnum>(*mode));
+		if (it != modes_.end())
+			currentMode_ = &it->second;
+		else
+			LOG(Awb, Error) << "Unsupported AWB mode " << *mode;
+	}
+}
+
+RGB<double> AwbBayes::gainsFromColourTemperature(double colourTemperature)
+{
+	/*
+	 * \todo: In the RaspberryPi code, the ct curve was interpolated in
+	 * the white point space (1/x) not in gains space. This feels counter
+	 * intuitive, as the gains are in linear space. But I can't prove it.
+	 */
+	const auto &gains = colourGainCurve_.getInterpolated(colourTemperature);
+	return { { gains[0], 1.0, gains[1] } };
+}
+
+AwbResult AwbBayes::calculateAwb(const AwbStats &stats, int lux)
+{
+	ipa::Pwl prior;
+	if (lux > 0) {
+		prior = priors_.getInterpolated(lux);
+		prior.map([](double x, double y) {
+			LOG(Awb, Debug) << "(" << x << "," << y << ")";
+		});
+	} else {
+		prior.append(0, 1.0);
+	}
+
+	double t = coarseSearch(prior, stats);
+	double r = ctR_.eval(t);
+	double b = ctB_.eval(t);
+	LOG(Awb, Debug)
+		<< "After coarse search: r " << r << " b " << b << " (gains r "
+		<< 1 / r << " b " << 1 / b << ")";
+
+	/*
+	 * Original comment from RaspberryPi:
+	 * Not entirely sure how to handle the fine search yet. Mostly the
+	 * estimated CT is already good enough, but the fine search allows us to
+	 * wander transversely off the CT curve. Under some illuminants, where
+	 * there may be more or less green light, this may prove beneficial,
+	 * though I probably need more real datasets before deciding exactly how
+	 * this should be controlled and tuned.
+	 */
+	fineSearch(t, r, b, prior, stats);
+	LOG(Awb, Debug)
+		<< "After fine search: r " << r << " b " << b << " (gains r "
+		<< 1 / r << " b " << 1 / b << ")";
+
+	return { { { 1.0 / r, 1.0, 1.0 / b } }, t };
+}
+
+double AwbBayes::coarseSearch(const ipa::Pwl &prior, const AwbStats &stats) const
+{
+	std::vector<Pwl::Point> points;
+	size_t bestPoint = 0;
+	double t = currentMode_->ctLo;
+	int spanR = -1;
+	int spanB = -1;
+	LimitsRecorder<double> errorLimits;
+	LimitsRecorder<double> priorLogLikelihoodLimits;
+
+	/* Step down the CT curve evaluating log likelihood. */
+	while (true) {
+		double r = ctR_.eval(t, &spanR);
+		double b = ctB_.eval(t, &spanB);
+		RGB<double> gains({ 1 / r, 1.0, 1 / b });
+		double delta2Sum = stats.computeColourError(gains);
+		double priorLogLikelihood = log(prior.eval(prior.domain().clamp(t)));
+		double finalLogLikelihood = delta2Sum - priorLogLikelihood;
+
+		errorLimits.record(delta2Sum);
+		priorLogLikelihoodLimits.record(priorLogLikelihood);
+
+		points.push_back({ { t, finalLogLikelihood } });
+		if (points.back().y() < points[bestPoint].y())
+			bestPoint = points.size() - 1;
+
+		if (t == currentMode_->ctHi)
+			break;
+
+		/*
+		 * Ensure even steps along the r/b curve by scaling them by the
+		 * current t.
+		 */
+		t = std::min(t + t / 10 * kSearchStep, currentMode_->ctHi);
+	}
+
+	t = points[bestPoint].x();
+	LOG(Awb, Debug) << "Coarse search found CT " << t
+			<< " error limits:" << errorLimits
+			<< " prior log likelihood limits:" << priorLogLikelihoodLimits;
+
+	/*
+	 * We have the best point of the search, but refine it with a quadratic
+	 * interpolation around its neighbors.
+	 */
+	if (points.size() > 2) {
+		bestPoint = std::clamp(bestPoint, std::size_t{ 1 }, points.size() - 2);
+		t = interpolateQuadratic(points[bestPoint - 1],
+					 points[bestPoint],
+					 points[bestPoint + 1]);
+		LOG(Awb, Debug)
+			<< "After quadratic refinement, coarse search has CT "
+			<< t;
+	}
+
+	return t;
+}
+
+void AwbBayes::fineSearch(double &t, double &r, double &b, ipa::Pwl const &prior, const AwbStats &stats) const
+{
+	int spanR = -1;
+	int spanB = -1;
+	double step = t / 10 * kSearchStep * 0.1;
+	int nsteps = 5;
+	ctR_.eval(t, &spanR);
+	ctB_.eval(t, &spanB);
+	double rDiff = ctR_.eval(t + nsteps * step, &spanR) -
+		       ctR_.eval(t - nsteps * step, &spanR);
+	double bDiff = ctB_.eval(t + nsteps * step, &spanB) -
+		       ctB_.eval(t - nsteps * step, &spanB);
+	Pwl::Point transverse({ bDiff, -rDiff });
+	if (transverse.length2() < 1e-6)
+		return;
+	/*
+	 * transverse is a unit vector orthogonal to the b vs. r function
+	 * (pointing outwards with r and b increasing)
+	 */
+	transverse = transverse / transverse.length();
+	double bestLogLikelihood = 0;
+	double bestT = 0;
+	Pwl::Point bestRB(0);
+	double transverseRange = transverseNeg_ + transversePos_;
+	const int maxNumDeltas = 12;
+	LimitsRecorder<double> errorLimits;
+	LimitsRecorder<double> priorLogLikelihoodLimits;
+
+
+	/* a transverse step approximately every 0.01 r/b units */
+	int numDeltas = floor(transverseRange * 100 + 0.5) + 1;
+	numDeltas = std::clamp(numDeltas, 3, maxNumDeltas);
+
+	/*
+	 * Step down CT curve. March a bit further if the transverse range is
+	 * large.
+	 */
+	nsteps += numDeltas;
+	for (int i = -nsteps; i <= nsteps; i++) {
+		double tTest = t + i * step;
+		double priorLogLikelihood =
+			log(prior.eval(prior.domain().clamp(tTest)));
+		priorLogLikelihoodLimits.record(priorLogLikelihood);
+		Pwl::Point rbStart{ { ctR_.eval(tTest, &spanR),
+				      ctB_.eval(tTest, &spanB) } };
+		Pwl::Point samples[maxNumDeltas];
+		int bestPoint = 0;
+
+		/*
+		 * Sample numDeltas points transversely *off* the CT curve
+		 * in the range [-transverseNeg, transversePos].
+		 * The x values of a sample contains the distance and the y
+		 * value contains the corresponding log likelihood.
+		 */
+		double transverseStep = transverseRange / (numDeltas - 1);
+		for (int j = 0; j < numDeltas; j++) {
+			auto &p = samples[j];
+			p.x() = -transverseNeg_ + transverseStep * j;
+			Pwl::Point rbTest = rbStart + transverse * p.x();
+			RGB<double> gains({ 1 / rbTest[0], 1.0, 1 / rbTest[1] });
+			double delta2Sum = stats.computeColourError(gains);
+			errorLimits.record(delta2Sum);
+			p.y() = delta2Sum - priorLogLikelihood;
+
+			if (p.y() < samples[bestPoint].y())
+				bestPoint = j;
+		}
+
+		/*
+		 * We have all samples transversely across the CT curve,
+		 * now let's do a quadratic interpolation for the best result.
+		 */
+		bestPoint = std::clamp(bestPoint, 1, numDeltas - 2);
+		double bestOffset = interpolateQuadratic(samples[bestPoint - 1],
+							 samples[bestPoint],
+							 samples[bestPoint + 1]);
+		Pwl::Point rbTest = rbStart + transverse * bestOffset;
+		RGB<double> gains({ 1 / rbTest[0], 1.0, 1 / rbTest[1] });
+		double delta2Sum = stats.computeColourError(gains);
+		errorLimits.record(delta2Sum);
+		double finalLogLikelihood = delta2Sum - priorLogLikelihood;
+
+		if (bestT == 0 || finalLogLikelihood < bestLogLikelihood) {
+			bestLogLikelihood = finalLogLikelihood;
+			bestT = tTest;
+			bestRB = rbTest;
+		}
+	}
+
+	t = bestT;
+	r = bestRB[0];
+	b = bestRB[1];
+	LOG(Awb, Debug)
+		<< "Fine search found t " << t << " r " << r << " b " << b
+		<< " error limits: " << errorLimits
+		<< " prior log likelihood limits: " << priorLogLikelihoodLimits;
+}
+
+/**
+ * \brief Find extremum of function
+ * \param[in] a First point
+ * \param[in] b Second point
+ * \param[in] c Third point
+ *
+ * Given 3 points on a curve, find the extremum of the function in that interval
+ * by fitting a quadratic.
+ *
+ * \return The x value of the extremum clamped to the interval [a.x(), c.x()]
+ */
+double AwbBayes::interpolateQuadratic(Pwl::Point const &a, Pwl::Point const &b,
+				      Pwl::Point const &c) const
+{
+	const double eps = 1e-3;
+	Pwl::Point ca = c - a;
+	Pwl::Point ba = b - a;
+	double denominator = 2 * (ba.y() * ca.x() - ca.y() * ba.x());
+	if (std::abs(denominator) > eps) {
+		double numerator = ba.y() * ca.x() * ca.x() - ca.y() * ba.x() * ba.x();
+		double result = numerator / denominator + a.x();
+		return std::max(a.x(), std::min(c.x(), result));
+	}
+	/* has degenerated to straight line segment */
+	return a.y() < c.y() - eps ? a.x() : (c.y() < a.y() - eps ? c.x() : b.x());
+}
+
+} /* namespace ipa */
+
+} /* namespace libcamera */