[PATCH v6 7/9] libtuning: Add initial AWB module

[Stefan Klug]

This AWB module 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.

Signed-off-by: Stefan Klug <stefan.klug at ideasonboard.com>
Reviewed-by: Kieran Bingham <kieran.bingham at ideasonboard.com>
Reviewed-by: Paul Elder <paul.elder at ideasonboard.com>
Reviewed-by: Laurent Pinchart <laurent.pinchart at ideasonboard.com>
---
 .../tuning/libtuning/modules/awb/__init__.py  |  6 ++++
 utils/tuning/libtuning/modules/awb/awb.py     | 36 +++++++++++++++++++
 utils/tuning/libtuning/modules/awb/rkisp1.py  | 27 ++++++++++++++
 3 files changed, 69 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..c154cf3b8609
--- /dev/null
+++ b/utils/tuning/libtuning/modules/awb/awb.py
@@ -0,0 +1,36 @@
+# 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.reshape(gains, (-1, 3))
+
+        return [{
+                    'ct': int(v[0]),
+                    'gains': [float(1.0 / v[1]), float(1.0 / v[2])]
+                } for v in gains]
diff --git a/utils/tuning/libtuning/modules/awb/rkisp1.py b/utils/tuning/libtuning/modules/awb/rkisp1.py
new file mode 100644
index 000000000000..0c95843b83d3
--- /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['colourGains'] = self.do_calculation(images)
+
+        return output
-- 
2.43.0