From 60d60c13672b81f2a57894467a569c6ba98ae895 Mon Sep 17 00:00:00 2001
From: Stefan Klug <stefan.klug@ideasonboard.com>
Date: Thu, 23 Jan 2025 12:40:58 +0100
Subject: libtuning: module: awb: Add bayes AWB support

To support the bayesian AWB algorithm in libtuning, the necessary data
needs to be collected and written to the tuning file.

Extend libtuning to calculate and output that additional data.

Prior probabilities and AwbModes are manually specified and not
calculated in the tuning process. Add sample values from the RaspberryPi
tuning files to the example config file.

Signed-off-by: Stefan Klug <stefan.klug@ideasonboard.com>
Reviewed-by: Paul Elder <paul.elder@ideasonboard.com>
Reviewed-by: Kieran Bingham <kieran.bingham@ideasonboard.com>
---
 utils/tuning/config-example.yaml             | 44 +++++++++++++++++++++++++++-
 utils/tuning/libtuning/modules/awb/awb.py    | 16 ++++++----
 utils/tuning/libtuning/modules/awb/rkisp1.py | 21 +++++++++----
 3 files changed, 68 insertions(+), 13 deletions(-)

diff --git a/utils/tuning/config-example.yaml b/utils/tuning/config-example.yaml
index 1b7f52cd..1bbb2757 100644
--- a/utils/tuning/config-example.yaml
+++ b/utils/tuning/config-example.yaml
@@ -5,7 +5,49 @@ general:
     do_alsc_colour: 1
     luminance_strength: 0.5
   awb:
-    greyworld: 0
+    # Algorithm can either be 'grey' or 'bayes'
+    algorithm: bayes
+    # Priors is only used for the bayes algorithm. They are defined in
+    # logarithmic space. A good staring point is:
+    # - lux: 0
+    #   ct: [ 2000, 3000, 13000 ]
+    #   probability: [ 1.0, 0.0, 0.0 ]
+    # - lux: 800
+    #   ct: [ 2000, 6000, 13000 ]
+    #   probability: [ 0.0, 2.0, 2.0 ]
+    # - lux: 1500
+    #   ct: [ 2000, 4000, 6000, 6500, 7000, 13000 ]
+    #   probability: [ 0.0, 1.0, 6.0, 7.0, 1.0, 1.0 ]
+    priors:
+      - lux: 0
+        ct: [ 2000, 13000 ]
+        probability: [ 0.0, 0.0 ]
+    AwbMode:
+      AwbAuto:
+        lo: 2500
+        hi: 8000
+      AwbIncandescent:
+        lo: 2500
+        hi: 3000
+      AwbTungsten:
+        lo: 3000
+        hi: 3500
+      AwbFluorescent:
+        lo: 4000
+        hi: 4700
+      AwbIndoor:
+        lo: 3000
+        hi: 5000
+      AwbDaylight:
+        lo: 5500
+        hi: 6500
+      AwbCloudy:
+        lo: 6500
+        hi: 8000
+      # One custom mode can be defined if needed
+      #AwbCustom:
+      #  lo: 2000
+      #  hi: 1300
   macbeth:
     small: 1
     show: 0
diff --git a/utils/tuning/libtuning/modules/awb/awb.py b/utils/tuning/libtuning/modules/awb/awb.py
index c154cf3b..0dc4f59d 100644
--- a/utils/tuning/libtuning/modules/awb/awb.py
+++ b/utils/tuning/libtuning/modules/awb/awb.py
@@ -27,10 +27,14 @@ class AWB(Module):
 
         imgs = [img for img in images if img.macbeth is not None]
 
-        gains, _, _ = awb(imgs, None, None, False)
-        gains = np.reshape(gains, (-1, 3))
+        ct_curve, transverse_pos, transverse_neg = awb(imgs, None, None, False)
+        ct_curve = np.reshape(ct_curve, (-1, 3))
+        gains = [{
+            'ct': int(v[0]),
+            'gains': [float(1.0 / v[1]), float(1.0 / v[2])]
+        } for v in ct_curve]
+
+        return {'colourGains': gains,
+                'transversePos': transverse_pos,
+                'transverseNeg': transverse_neg}
 
-        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
index 0c95843b..d562d26e 100644
--- a/utils/tuning/libtuning/modules/awb/rkisp1.py
+++ b/utils/tuning/libtuning/modules/awb/rkisp1.py
@@ -6,9 +6,6 @@
 
 from .awb import AWB
 
-import libtuning as lt
-
-
 class AWBRkISP1(AWB):
     hr_name = 'AWB (RkISP1)'
     out_name = 'Awb'
@@ -20,8 +17,20 @@ class AWBRkISP1(AWB):
         return True
 
     def process(self, config: dict, images: list, outputs: dict) -> dict:
-        output = {}
-
-        output['colourGains'] = self.do_calculation(images)
+        if not 'awb' in config['general']:
+            raise ValueError('AWB configuration missing')
+        awb_config = config['general']['awb']
+        algorithm = awb_config['algorithm']
+
+        output = {'algorithm': algorithm}
+        data = self.do_calculation(images)
+        if algorithm == 'grey':
+            output['colourGains'] = data['colourGains']
+        elif algorithm == 'bayes':
+            output['AwbMode'] = awb_config['AwbMode']
+            output['priors'] = awb_config['priors']
+            output.update(data)
+        else:
+            raise ValueError(f"Unknown AWB algorithm {output['algorithm']}")
 
         return output
-- 
cgit v1.2.1