libtuning: Add initial AWB module

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@ideasonboard.com>
Reviewed-by: Kieran Bingham <kieran.bingham@ideasonboard.com>
Reviewed-by: Paul Elder <paul.elder@ideasonboard.com>
Reviewed-by: Laurent Pinchart <laurent.pinchart@ideasonboard.com>

3 files changed, 69 insertions, 0 deletions

diff --git a/utils/tuning/libtuning/modules/awb/__init__.py b/utils/tuning/libtuning/modules/awb/__init__.py
new file mode 100644
index 00000000..2d67f10c
--- /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 00000000..c154cf3b
--- /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 00000000..0c95843b
--- /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