From patchwork Tue Jan 27 17:13:17 2026 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: David Plowman X-Patchwork-Id: 25980 Return-Path: X-Original-To: parsemail@patchwork.libcamera.org Delivered-To: parsemail@patchwork.libcamera.org Received: from lancelot.ideasonboard.com (lancelot.ideasonboard.com [92.243.16.209]) by patchwork.libcamera.org (Postfix) with ESMTPS id 261F7C328D for ; Tue, 27 Jan 2026 17:15:20 +0000 (UTC) Received: from lancelot.ideasonboard.com (localhost [IPv6:::1]) by lancelot.ideasonboard.com (Postfix) with ESMTP id 85A9F61FD6; Tue, 27 Jan 2026 18:15:18 +0100 (CET) Authentication-Results: lancelot.ideasonboard.com; dkim=pass (2048-bit key; unprotected) header.d=raspberrypi.com header.i=@raspberrypi.com header.b="OucjrZTA"; dkim-atps=neutral Received: from mail-wm1-x32d.google.com (mail-wm1-x32d.google.com [IPv6:2a00:1450:4864:20::32d]) by lancelot.ideasonboard.com (Postfix) with ESMTPS id 88E3361FD2 for ; Tue, 27 Jan 2026 18:15:13 +0100 (CET) Received: by mail-wm1-x32d.google.com with SMTP id 5b1f17b1804b1-47ee3a63300so67258345e9.2 for ; Tue, 27 Jan 2026 09:15:13 -0800 (PST) DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=raspberrypi.com; s=google; t=1769534113; x=1770138913; darn=lists.libcamera.org; h=content-transfer-encoding:mime-version:references:in-reply-to :message-id:date:subject:cc:to:from:from:to:cc:subject:date :message-id:reply-to; bh=x2bQHuSs/sNthSP3D+231dZk6wtL/itXkk2+f/mLY6w=; b=OucjrZTAuZQ4PeWU4bpoutZ4PvnxZ09ov1mutzkKJqxgOMfJUsF3gYiFTbn1QAa/pS Ag77Jfnkrz+1kDmJWnXU9yorZzb9BycNQwCS1falwDJ1PFiq6H7GCt7JOAOQGxhMj5Hi 0Rr72HG9DCrE8xxVYuelq+mIhzNaVUFsNoye87jN/uVSbQ4Q5EABEkWkzV/DUhuIdiIn H+3Qr8eTgrpFYnZ017qRVFpm1D39OwpKCjEanbAiv0AsBIDMdksoWxbZ1h4QemqWd3uX fp+1c3xbbckZCHXSYRNkbX/YQrL510kKt9+DU0lo8afpHgso/JSGjKUNR+xtr2UNKxs8 9V+Q== X-Google-DKIM-Signature: v=1; a=rsa-sha256; c=relaxed/relaxed; d=1e100.net; s=20230601; t=1769534113; x=1770138913; h=content-transfer-encoding:mime-version:references:in-reply-to :message-id:date:subject:cc:to:from:x-gm-gg:x-gm-message-state:from :to:cc:subject:date:message-id:reply-to; bh=x2bQHuSs/sNthSP3D+231dZk6wtL/itXkk2+f/mLY6w=; b=chBxjjUmJASro1frzsZcrQ879TEShlBDC4pv9vUyAd/gUeDKVci0hQwr01akSnq4bV cG2JC2FR/OdokzxsQuOOpd7TBncVAwZLmtktpRL+5GVqWfWVzuSeNRFPt0QWGIDQcW76 X2hLDa5PG7KP86Qp2s97k+926Bmws2C/TCjrMMDZlUQjH7/StuUUQvA97OEkumsVaW46 ZgCbwgziwYdSTYg/hY+HvukyQqI40QnS3spoWhC3fAsykKGOiULI9ydy+y6tM3jZtvUH 8mnIufgGg0jHLPRmwf3jf8Ev7hpHkDBQXAKyjgjkDAbYOpJGlt79M9OW1xumayh3vrsp zZGw== X-Gm-Message-State: AOJu0YwnsawdEIqyuwkCEfltvG7OmKBIUFWXtIWl5doH+6x9UfuuhnIs jGlNf/B0MM9y6TY7rj4QL6dbyl2vf5w11w1uKmINk0/T5pYWVi70Ypb7UnEsmlzzNDOsgrd5CvZ HtCRd X-Gm-Gg: AZuq6aLlMjMzSRmst6hIXTm1Nxdb+o2kCwcuDOLEqqGp1S9gJZLjQostmnROUlOTFS3 6Mwr3ljE4JqcNQbXFr1EWR3OGPRy7RWI48nYC7vdWZPEaJTvuumf4iKWE5fd1UK7gQrpBYDoNrr 37Cy24FbsntOqUjGXVQfreMT2s9izLm1rQxk2+pYv5NbmWeceGSCqXQCpQdgHcDzbKHNh8Pq3wS JPGuPFqSGkwqTeuom6WvZ7RUmAms5F6G5a8g1Nb9JeoYrv7r9vO7lD4ytbB9LB8eQjQVbcBnqyq xPjcCwfyY72468BMKu3oNpINRAjauVKZf4xKG0h5OMiDgglMGRCtSAvx0cmLobe8We+pcdvoZit +v9vrUNTwjnQBK9keh0vFFhSvC6xdGHfyqBAxrtYoGzrCFvZyLqURBSCKClMldgE1wzSeYvmLSb pBUFx1GM5b18EpVEYsr4z52cFROJFaLknnfY0Fc2q40w0DaVAPPPN3lvM/e9eaELR1S/Uj67/3F iynb7O1O5qXiVhpwSdtC2kcvblvhQ== X-Received: by 2002:a05:600c:4e54:b0:46e:6d5f:f68 with SMTP id 5b1f17b1804b1-4806c7cfb14mr13159595e9.12.1769534112411; Tue, 27 Jan 2026 09:15:12 -0800 (PST) Received: from davidp-pi5.pitowers.org ([2a00:1098:3142:1f:88ea:c658:5b20:5e46]) by smtp.gmail.com with ESMTPSA id 5b1f17b1804b1-4806cddffe9sm12267865e9.4.2026.01.27.09.15.11 (version=TLS1_3 cipher=TLS_AES_256_GCM_SHA384 bits=256/256); Tue, 27 Jan 2026 09:15:11 -0800 (PST) From: David Plowman To: libcamera-devel@lists.libcamera.org Cc: Peter Bailey , David Plowman , Naushir Patuck Subject: [PATCH v6 2/4] ipa: rpi: controller: awb: Add Neural Network AWB Date: Tue, 27 Jan 2026 17:13:17 +0000 Message-ID: <20260127171506.17480-3-david.plowman@raspberrypi.com> X-Mailer: git-send-email 2.47.3 In-Reply-To: <20260127171506.17480-1-david.plowman@raspberrypi.com> References: <20260127171506.17480-1-david.plowman@raspberrypi.com> MIME-Version: 1.0 X-BeenThere: libcamera-devel@lists.libcamera.org X-Mailman-Version: 2.1.29 Precedence: list List-Id: List-Unsubscribe: , List-Archive: List-Post: List-Help: List-Subscribe: , Errors-To: libcamera-devel-bounces@lists.libcamera.org Sender: "libcamera-devel" From: Peter Bailey Add an AWB algorithm which uses neural networks. Signed-off-by: Peter Bailey Reviewed-by: David Plowman Reviewed-by: Naushir Patuck --- meson_options.txt | 5 + src/ipa/rpi/controller/meson.build | 9 + src/ipa/rpi/controller/rpi/awb_nn.cpp | 437 ++++++++++++++++++++++++++ 3 files changed, 451 insertions(+) create mode 100644 src/ipa/rpi/controller/rpi/awb_nn.cpp diff --git a/meson_options.txt b/meson_options.txt index c052e85a..07847294 100644 --- a/meson_options.txt +++ b/meson_options.txt @@ -76,6 +76,11 @@ option('qcam', value : 'auto', description : 'Compile the qcam test application') +option('rpi-awb-nn', + type : 'feature', + value : 'auto', + description : 'Enable the Raspberry Pi Neural Network AWB algorithm') + option('test', type : 'boolean', value : false, diff --git a/src/ipa/rpi/controller/meson.build b/src/ipa/rpi/controller/meson.build index c8637906..03ee7c20 100644 --- a/src/ipa/rpi/controller/meson.build +++ b/src/ipa/rpi/controller/meson.build @@ -32,6 +32,15 @@ rpi_ipa_controller_deps = [ libcamera_private, ] +tflite_dep = dependency('tensorflow-lite', required : get_option('rpi-awb-nn')) + +if tflite_dep.found() + rpi_ipa_controller_sources += files([ + 'rpi/awb_nn.cpp', + ]) + rpi_ipa_controller_deps += tflite_dep +endif + rpi_ipa_controller_lib = static_library('rpi_ipa_controller', rpi_ipa_controller_sources, include_directories : libipa_includes, dependencies : rpi_ipa_controller_deps) diff --git a/src/ipa/rpi/controller/rpi/awb_nn.cpp b/src/ipa/rpi/controller/rpi/awb_nn.cpp new file mode 100644 index 00000000..1c205126 --- /dev/null +++ b/src/ipa/rpi/controller/rpi/awb_nn.cpp @@ -0,0 +1,437 @@ +/* SPDX-License-Identifier: BSD-2-Clause */ +/* + * Copyright (C) 2025, Raspberry Pi Ltd + * + * AWB control algorithm using neural network + * + * The AWB Neural Network algorithm can be run entirely with the code here + * and the suppllied TFLite models. Those interested in the full model + * definitions, or who may want to re-train the models should visit + * + * https://github.com/raspberrypi/awb_nn + * + * where you will find full source code for the models, the full datasets + * used for training our supplied models, and full instructions for capturing + * your own images and re-training the models for your own use cases. + */ + +#include + +#include +#include + +#include +#include +#include + +#include "../awb_algorithm.h" +#include "../awb_status.h" +#include "../lux_status.h" +#include "libipa/pwl.h" + +#include "awb.h" + +using namespace libcamera; + +LOG_DECLARE_CATEGORY(RPiAwb) + +constexpr double kDefaultCT = 4500.0; + +/* + * The neural networks are trained to work on images rendered at a canonical + * colour temperature. That value is 5000K, which must be reproduced here. + */ +constexpr double kNetworkCanonicalCT = 5000.0; + +#define NAME "rpi.nn.awb" + +namespace RPiController { + +struct AwbNNConfig { + AwbNNConfig() = default; + int read(const libcamera::YamlObject ¶ms, AwbConfig &config); + + /* An empty model will check default locations for model.tflite */ + std::string model; + float minTemp; + float maxTemp; + + bool enableNn; + + /* CCM matrix for canonical network CT */ + double ccm[9]; +}; + +class AwbNN : public Awb +{ +public: + AwbNN(Controller *controller = nullptr); + ~AwbNN(); + char const *name() const override; + void initialise() override; + int read(const libcamera::YamlObject ¶ms) override; + +protected: + void doAwb() override; + void prepareStats() override; + +private: + bool isAutoEnabled() const; + AwbNNConfig nnConfig_; + void transverseSearch(double t, double &r, double &b); + RGB processZone(RGB zone, float red_gain, float blue_gain); + void awbNN(); + void loadModel(); + + libcamera::Size zoneSize_; + std::unique_ptr model_; + std::unique_ptr interpreter_; +}; + +int AwbNNConfig::read(const libcamera::YamlObject ¶ms, AwbConfig &config) +{ + model = params["model"].get(""); + minTemp = params["min_temp"].get(2800.0); + maxTemp = params["max_temp"].get(7600.0); + + for (int i = 0; i < 9; i++) + ccm[i] = params["ccm"][i].get(0.0); + + enableNn = params["enable_nn"].get(1); + + if (enableNn) { + if (!config.hasCtCurve()) { + LOG(RPiAwb, Error) << "CT curve not specified"; + enableNn = false; + } + + if (!model.empty() && model.find(".tflite") == std::string::npos) { + LOG(RPiAwb, Error) << "Model must be a .tflite file"; + enableNn = false; + } + + bool validCcm = true; + for (int i = 0; i < 9; i++) + if (ccm[i] == 0.0) + validCcm = false; + + if (!validCcm) { + LOG(RPiAwb, Error) << "CCM not specified or invalid"; + enableNn = false; + } + + if (!enableNn) { + LOG(RPiAwb, Warning) << "Neural Network AWB misconfigured - switch to Grey method"; + } + } + + if (!enableNn) { + config.sensitivityR = config.sensitivityB = 1.0; + config.greyWorld = true; + } + + return 0; +} + +AwbNN::AwbNN(Controller *controller) + : Awb(controller) +{ + zoneSize_ = getHardwareConfig().awbRegions; +} + +AwbNN::~AwbNN() +{ +} + +char const *AwbNN::name() const +{ + return NAME; +} + +int AwbNN::read(const libcamera::YamlObject ¶ms) +{ + int ret; + + ret = config_.read(params); + if (ret) + return ret; + + ret = nnConfig_.read(params, config_); + if (ret) + return ret; + + return 0; +} + +static bool checkTensorShape(TfLiteTensor *tensor, const int *expectedDims, const int expectedDimsSize) +{ + return std::equal(expectedDims, expectedDims + expectedDimsSize, + tensor->dims->data, tensor->dims->data + tensor->dims->size); +} + +static std::string buildDimString(const int *dims, const int dimsSize) +{ + return "[" + utils::join(Span(dims, dimsSize), ",") + "]"; +} + +void AwbNN::loadModel() +{ + std::string modelPath; + if (getTarget() == "bcm2835") { + modelPath = "/ipa/rpi/vc4/awb_model.tflite"; + } else { + modelPath = "/ipa/rpi/pisp/awb_model.tflite"; + } + + if (nnConfig_.model.empty()) { + std::string root = utils::libcameraSourcePath(); + if (!root.empty()) { + modelPath = root + modelPath; + } else { + modelPath = LIBCAMERA_DATA_DIR + modelPath; + } + + if (!File::exists(modelPath)) { + LOG(RPiAwb, Error) << "No model file found in standard locations"; + nnConfig_.enableNn = false; + return; + } + } else { + modelPath = nnConfig_.model; + } + + LOG(RPiAwb, Debug) << "Attempting to load model from: " << modelPath; + + model_ = tflite::FlatBufferModel::BuildFromFile(modelPath.c_str()); + + if (!model_) { + LOG(RPiAwb, Error) << "Failed to load model from " << modelPath; + nnConfig_.enableNn = false; + return; + } + + tflite::MutableOpResolver resolver; + tflite::ops::builtin::BuiltinOpResolver builtin_resolver; + resolver.AddAll(builtin_resolver); + tflite::InterpreterBuilder(*model_, resolver)(&interpreter_); + if (!interpreter_) { + LOG(RPiAwb, Error) << "Failed to build interpreter for model " << nnConfig_.model; + nnConfig_.enableNn = false; + return; + } + + interpreter_->AllocateTensors(); + TfLiteTensor *inputTensor = interpreter_->input_tensor(0); + TfLiteTensor *inputLuxTensor = interpreter_->input_tensor(1); + TfLiteTensor *outputTensor = interpreter_->output_tensor(0); + if (!inputTensor || !inputLuxTensor || !outputTensor) { + LOG(RPiAwb, Error) << "Model missing input or output tensor"; + nnConfig_.enableNn = false; + return; + } + + const int expectedInputDims[] = { 1, (int)zoneSize_.height, (int)zoneSize_.width, 3 }; + const int expectedInputLuxDims[] = { 1 }; + const int expectedOutputDims[] = { 1 }; + + if (!checkTensorShape(inputTensor, expectedInputDims, 4)) { + LOG(RPiAwb, Error) << "Model input tensor dimension mismatch. Expected: " << buildDimString(expectedInputDims, 4) + << ", Got: " << buildDimString(inputTensor->dims->data, inputTensor->dims->size); + nnConfig_.enableNn = false; + return; + } + + if (!checkTensorShape(inputLuxTensor, expectedInputLuxDims, 1)) { + LOG(RPiAwb, Error) << "Model input lux tensor dimension mismatch. Expected: " << buildDimString(expectedInputLuxDims, 1) + << ", Got: " << buildDimString(inputLuxTensor->dims->data, inputLuxTensor->dims->size); + nnConfig_.enableNn = false; + return; + } + + if (!checkTensorShape(outputTensor, expectedOutputDims, 1)) { + LOG(RPiAwb, Error) << "Model output tensor dimension mismatch. Expected: " << buildDimString(expectedOutputDims, 1) + << ", Got: " << buildDimString(outputTensor->dims->data, outputTensor->dims->size); + nnConfig_.enableNn = false; + return; + } + + if (inputTensor->type != kTfLiteFloat32 || inputLuxTensor->type != kTfLiteFloat32 || outputTensor->type != kTfLiteFloat32) { + LOG(RPiAwb, Error) << "Model input and output tensors must be float32"; + nnConfig_.enableNn = false; + return; + } + + LOG(RPiAwb, Info) << "Model loaded successfully from " << modelPath; + LOG(RPiAwb, Debug) << "Model validation successful - Input Image: " + << buildDimString(expectedInputDims, 4) + << ", Input Lux: " << buildDimString(expectedInputLuxDims, 1) + << ", Output: " << buildDimString(expectedOutputDims, 1) << " floats"; +} + +void AwbNN::initialise() +{ + Awb::initialise(); + + if (nnConfig_.enableNn) { + loadModel(); + if (!nnConfig_.enableNn) { + LOG(RPiAwb, Warning) << "Neural Network AWB failed to load - switch to Grey method"; + config_.greyWorld = true; + config_.sensitivityR = config_.sensitivityB = 1.0; + } + } +} + +void AwbNN::prepareStats() +{ + zones_.clear(); + /* + * LSC has already been applied to the stats in this pipeline, so stop + * any LSC compensation. We also ignore config_.fast in this version. + */ + generateStats(zones_, statistics_, 0.0, 0.0, getGlobalMetadata(), 0.0, 0.0, 0.0); + /* + * apply sensitivities, so values appear to come from our "canonical" + * sensor. + */ + for (auto &zone : zones_) { + zone.R *= config_.sensitivityR; + zone.B *= config_.sensitivityB; + } +} + +void AwbNN::transverseSearch(double t, double &r, double &b) +{ + int spanR = -1, spanB = -1; + config_.ctR.eval(t, &spanR); + config_.ctB.eval(t, &spanB); + + const int diff = 10; + double rDiff = config_.ctR.eval(t + diff, &spanR) - + config_.ctR.eval(t - diff, &spanR); + double bDiff = config_.ctB.eval(t + diff, &spanB) - + config_.ctB.eval(t - diff, &spanB); + + ipa::Pwl::Point transverse({ bDiff, -rDiff }); + if (transverse.length2() < 1e-6) + return; + + transverse = transverse / transverse.length(); + double transverseRange = config_.transverseNeg + config_.transversePos; + const int maxNumDeltas = 12; + int numDeltas = floor(transverseRange * 100 + 0.5) + 1; + numDeltas = std::clamp(numDeltas, 3, maxNumDeltas); + + ipa::Pwl::Point points[maxNumDeltas]; + int bestPoint = 0; + + for (int i = 0; i < numDeltas; i++) { + points[i][0] = -config_.transverseNeg + + (transverseRange * i) / (numDeltas - 1); + ipa::Pwl::Point rbTest = ipa::Pwl::Point({ r, b }) + + transverse * points[i].x(); + double rTest = rbTest.x(), bTest = rbTest.y(); + double gainR = 1 / rTest, gainB = 1 / bTest; + double delta2Sum = computeDelta2Sum(gainR, gainB, 0.0, 0.0); + points[i][1] = delta2Sum; + if (points[i].y() < points[bestPoint].y()) + bestPoint = i; + } + + bestPoint = std::clamp(bestPoint, 1, numDeltas - 2); + ipa::Pwl::Point rbBest = ipa::Pwl::Point({ r, b }) + + transverse * interpolateQuadatric(points[bestPoint - 1], + points[bestPoint], + points[bestPoint + 1]); + r = rbBest.x(); + b = rbBest.y(); +} + +AwbNN::RGB AwbNN::processZone(AwbNN::RGB zone, float redGain, float blueGain) +{ + /* + * Renders the pixel at canonical network colour temperature + */ + RGB zoneGains = zone; + + zoneGains.R *= redGain; + zoneGains.G *= 1.0; + zoneGains.B *= blueGain; + + RGB zoneCcm; + + zoneCcm.R = nnConfig_.ccm[0] * zoneGains.R + nnConfig_.ccm[1] * zoneGains.G + nnConfig_.ccm[2] * zoneGains.B; + zoneCcm.G = nnConfig_.ccm[3] * zoneGains.R + nnConfig_.ccm[4] * zoneGains.G + nnConfig_.ccm[5] * zoneGains.B; + zoneCcm.B = nnConfig_.ccm[6] * zoneGains.R + nnConfig_.ccm[7] * zoneGains.G + nnConfig_.ccm[8] * zoneGains.B; + + return zoneCcm; +} + +void AwbNN::awbNN() +{ + float *inputData = interpreter_->typed_input_tensor(0); + float *inputLux = interpreter_->typed_input_tensor(1); + + float redGain = 1.0 / config_.ctR.eval(kNetworkCanonicalCT); + float blueGain = 1.0 / config_.ctB.eval(kNetworkCanonicalCT); + + for (unsigned int i = 0; i < zoneSize_.height; i++) { + for (unsigned int j = 0; j < zoneSize_.width; j++) { + unsigned int zoneIdx = i * zoneSize_.width + j; + + RGB processedZone = processZone(zones_[zoneIdx] * (1.0 / 65535), redGain, blueGain); + unsigned int baseIdx = zoneIdx * 3; + + inputData[baseIdx + 0] = static_cast(processedZone.R); + inputData[baseIdx + 1] = static_cast(processedZone.G); + inputData[baseIdx + 2] = static_cast(processedZone.B); + } + } + + inputLux[0] = static_cast(lux_); + + TfLiteStatus status = interpreter_->Invoke(); + if (status != kTfLiteOk) { + LOG(RPiAwb, Error) << "Model inference failed with status: " << status; + return; + } + + float *outputData = interpreter_->typed_output_tensor(0); + + double t = outputData[0]; + + LOG(RPiAwb, Debug) << "Model output temperature: " << t; + + t = std::clamp(t, mode_->ctLo, mode_->ctHi); + + double r = config_.ctR.eval(t); + double b = config_.ctB.eval(t); + + transverseSearch(t, r, b); + + LOG(RPiAwb, Debug) << "After transverse search: Temperature: " << t << " Red gain: " << 1.0 / r << " Blue gain: " << 1.0 / b; + + asyncResults_.temperatureK = t; + asyncResults_.gainR = 1.0 / r * config_.sensitivityR; + asyncResults_.gainG = 1.0; + asyncResults_.gainB = 1.0 / b * config_.sensitivityB; +} + +void AwbNN::doAwb() +{ + prepareStats(); + if (zones_.size() == (zoneSize_.width * zoneSize_.height) && nnConfig_.enableNn) + awbNN(); + else + awbGrey(); + statistics_.reset(); +} + +/* Register algorithm with the system. */ +static Algorithm *create(Controller *controller) +{ + return new AwbNN(controller); +} +static RegisterAlgorithm reg(NAME, &create); + +} /* namespace RPiController */