From 81fd85c2711787167e45bcdd67f7986c99602286 Mon Sep 17 00:00:00 2001
From: Naushir Patuck <naush@raspberrypi.com>
Date: Mon, 9 Sep 2024 13:37:07 +0100
Subject: [PATCH] ipa: rpi: Handle the new CNN controls in the IPA

Add code to handle the new CNN vendor controls in the Raspberry Pi IPA.

The value of CnnInputTensorInfo is cached as it is the only stateful
input control.

All other controls are output controls, and the values are copied into
directly from the rpiMetadata object if present. The camera helpers
populate the rpiMetadata object if the sensor supports on-board CNN
processing, such as the IMX500.

Signed-off-by: Naushir Patuck <naush@raspberrypi.com>
---
 src/ipa/rpi/common/ipa_base.cpp | 52 ++++++++++++++++++++++++++++++++-
 src/ipa/rpi/common/ipa_base.h   |  1 +
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/src/ipa/rpi/common/ipa_base.cpp b/src/ipa/rpi/common/ipa_base.cpp
index faa77719..b2baa52f 100644
--- a/src/ipa/rpi/common/ipa_base.cpp
+++ b/src/ipa/rpi/common/ipa_base.cpp
@@ -82,6 +82,7 @@ const ControlInfoMap::Map ipaControls{
 						static_cast<int64_t>(defaultMinFrameDuration.get<std::micro>()) } }) },
 	{ &controls::draft::NoiseReductionMode, ControlInfo(controls::draft::NoiseReductionModeValues) },
 	{ &controls::rpi::StatsOutputEnable, ControlInfo(false, true, false) },
+	{ &controls::rpi::CnnEnableInputTensor, ControlInfo(false, true, false) },
 };
 
 /* IPA controls handled conditionally, if the sensor is not mono */
@@ -123,7 +124,7 @@ namespace ipa::RPi {
 IpaBase::IpaBase()
 	: controller_(), frameLengths_(FrameLengthsQueueSize, 0s), statsMetadataOutput_(false),
 	  stitchSwapBuffers_(false), frameCount_(0), mistrustCount_(0), lastRunTimestamp_(0),
-	  firstStart_(true), flickerState_({ 0, 0s }), awbEnabled_(true)
+	  firstStart_(true), flickerState_({ 0, 0s }), cnnEnableInputTensor_(false), awbEnabled_(true)
 {
 }
 
@@ -1456,6 +1457,10 @@ void IpaBase::applyControls(const ControlList &controls)
 			statsMetadataOutput_ = ctrl.second.get<bool>();
 			break;
 
+		case controls::rpi::CNN_ENABLE_INPUT_TENSOR:
+			cnnEnableInputTensor_ = ctrl.second.get<bool>();
+			break;
+
 		default:
 			LOG(IPARPI, Warning)
 				<< "Ctrl " << controls::controls.at(ctrl.first)->name()
@@ -1645,6 +1650,51 @@ void IpaBase::reportMetadata(unsigned int ipaContext)
 			libcameraMetadata_.set(controls::HdrChannel, controls::HdrChannelNone);
 	}
 
+	const std::shared_ptr<uint8_t[]> *inputTensor =
+		rpiMetadata.getLocked<std::shared_ptr<uint8_t[]>>("cnn.input_tensor");
+	if (cnnEnableInputTensor_ && inputTensor) {
+		unsigned int size = *rpiMetadata.getLocked<unsigned int>("cnn.input_tensor_size");
+		Span<const uint8_t> tensor{ inputTensor->get(), size };
+		libcameraMetadata_.set(controls::rpi::CnnInputTensor, tensor);
+		/* No need to keep these big buffers any more. */
+		rpiMetadata.eraseLocked("cnn.input_tensor");
+	}
+
+	const RPiController::CnnInputTensorInfo *inputTensorInfo =
+		rpiMetadata.getLocked<RPiController::CnnInputTensorInfo>("cnn.input_tensor_info");
+	if (inputTensorInfo) {
+		Span<const uint8_t> tensorInfo{ reinterpret_cast<const uint8_t *>(inputTensorInfo),
+						sizeof(*inputTensorInfo) };
+		libcameraMetadata_.set(controls::rpi::CnnInputTensorInfo, tensorInfo);
+	}
+
+	const std::shared_ptr<float[]> *outputTensor =
+		rpiMetadata.getLocked<std::shared_ptr<float[]>>("cnn.output_tensor");
+	if (outputTensor) {
+		unsigned int size = *rpiMetadata.getLocked<unsigned int>("cnn.output_tensor_size");
+		Span<const float> tensor{ reinterpret_cast<const float *>(outputTensor->get()),
+					  size };
+		libcameraMetadata_.set(controls::rpi::CnnOutputTensor, tensor);
+		/* No need to keep these big buffers any more. */
+		rpiMetadata.eraseLocked("cnn.output_tensor");
+	}
+
+	const RPiController::CnnOutputTensorInfo *outputTensorInfo =
+		rpiMetadata.getLocked<RPiController::CnnOutputTensorInfo>("cnn.output_tensor_info");
+	if (outputTensorInfo) {
+		Span<const uint8_t> tensorInfo{ reinterpret_cast<const uint8_t *>(outputTensorInfo),
+						sizeof(*outputTensorInfo) };
+		libcameraMetadata_.set(controls::rpi::CnnOutputTensorInfo, tensorInfo);
+	}
+
+	const RPiController::CnnKpiInfo *kpiInfo =
+		rpiMetadata.getLocked<RPiController::CnnKpiInfo>("cnn.kpi_info");
+	if (kpiInfo) {
+		libcameraMetadata_.set(controls::rpi::CnnKpiInfo,
+				       { static_cast<int32_t>(kpiInfo->dnnRuntime),
+					 static_cast<int32_t>(kpiInfo->dspRuntime) });
+	}
+
 	metadataReady.emit(libcameraMetadata_);
 }
 
diff --git a/src/ipa/rpi/common/ipa_base.h b/src/ipa/rpi/common/ipa_base.h
index 90f018b2..90e21756 100644
--- a/src/ipa/rpi/common/ipa_base.h
+++ b/src/ipa/rpi/common/ipa_base.h
@@ -141,6 +141,7 @@ private:
 		utils::Duration manualPeriod;
 	} flickerState_;
 
+	bool cnnEnableInputTensor_;
 	bool awbEnabled_;
 
 	utils::Duration controllerMinFrameDuration_;