1 files changed, 15 insertions, 6 deletions
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