[PATCH v1 5/6] libtuning: Add initial AWB module

[Stefan Klug]

This AWB modules uses the awb function from Raspberry Pi to calculate
the needed white balance gains per colour temperature. It stores these
gains in the tuning file. Currently they are only used for the manual
colour temperature control. It is likely that they will be used with
more complex awb algorithms.

Signed-off-by: Stefan Klug <stefan.klug at ideasonboard.com>
---
 .../tuning/libtuning/modules/awb/__init__.py  |  6 +++
 utils/tuning/libtuning/modules/awb/awb.py     | 42 +++++++++++++++++++
 utils/tuning/libtuning/modules/awb/rkisp1.py  | 27 ++++++++++++
 3 files changed, 75 insertions(+)
 create mode 100644 utils/tuning/libtuning/modules/awb/__init__.py
 create mode 100644 utils/tuning/libtuning/modules/awb/awb.py
 create mode 100644 utils/tuning/libtuning/modules/awb/rkisp1.py

diff --git a/utils/tuning/libtuning/modules/awb/__init__.py b/utils/tuning/libtuning/modules/awb/__init__.py
new file mode 100644
index 000000000000..2d67f10cfc4f
--- /dev/null
+++ b/utils/tuning/libtuning/modules/awb/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+#
+# Copyright (C) 2024, Ideas On Board
+
+from libtuning.modules.awb.awb import AWB
+from libtuning.modules.awb.rkisp1 import AWBRkISP1
diff --git a/utils/tuning/libtuning/modules/awb/awb.py b/utils/tuning/libtuning/modules/awb/awb.py
new file mode 100644
index 000000000000..5680a44fd3e3
--- /dev/null
+++ b/utils/tuning/libtuning/modules/awb/awb.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+#
+# Copyright (C) 2024, Ideas On Board
+
+import logging
+
+from ..module import Module
+
+from libtuning.ctt_awb import awb
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class AWB(Module):
+    type = 'awb'
+    hr_name = 'AWB (Base)'
+    out_name = 'GenericAWB'
+
+    def __init__(self, *,
+                 debug: list):
+        super().__init__()
+
+        self.debug = debug
+
+    def do_calculation(self, images):
+        logger.info('Starting AWB calculation')
+
+        imgs = [img for img in images if img.macbeth is not None]
+
+        gains, _, _ = awb(imgs, None, None, False)
+        gains = np.array(gains)
+        gains = gains.reshape(-1, 3)
+
+        res = []
+        for v in gains:
+            res.append({
+                'ct': int(v[0]),
+                'gains': [float(1.0 / v[1]), float(1.0 / v[2])]
+            })
+
+        return res
diff --git a/utils/tuning/libtuning/modules/awb/rkisp1.py b/utils/tuning/libtuning/modules/awb/rkisp1.py
new file mode 100644
index 000000000000..e3ddeb50c9a1
--- /dev/null
+++ b/utils/tuning/libtuning/modules/awb/rkisp1.py
@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+#
+# Copyright (C) 2024, Ideas On Board
+#
+# AWB module for tuning rkisp1
+
+from .awb import AWB
+
+import libtuning as lt
+
+
+class AWBRkISP1(AWB):
+    hr_name = 'AWB (RkISP1)'
+    out_name = 'Awb'
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def validate_config(self, config: dict) -> bool:
+        return True
+
+    def process(self, config: dict, images: list, outputs: dict) -> dict:
+        output = {}
+
+        output['gains'] = self.do_calculation(images)
+
+        return output
-- 
2.43.0