risp4ml/isp_stages/wbg.c

/*
 * Copyright 2023 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "risp4ml/isp_stages/wbg.h"

#include "risp4ml/common/utils.h"

#define MAX(a, b) (((a) > (b)) ? (a) : (b))

static const uint16_t kWbgFractional = kRawPipelineFraction;
static const uint16_t kWbgUnityGain = 1 << kWbgFractional;
static WbgParams wbg_params = {
    .enable = true,
    .fixed = false,
    .gains = {kWbgUnityGain, kWbgUnityGain, kWbgUnityGain, kWbgUnityGain}};

void set_wbg_params(WbgParams* params) { wbg_params = *params; }

static void compute_wbg_gain(Image* img) {
  // Calculate the white-balance gain values using the "gray world" algorithm
  int64_t sum_of_reds = 0;
  uint32_t num_of_reds = 0;
  // will use only one of the greens for scaling, since the difference between
  // the two green sensor pixels is negligible
  int64_t sum_of_greens = 0;
  uint32_t num_of_greens = 0;
  int64_t sum_of_blues = 0;
  uint32_t num_of_blues = 0;

  for (uint16_t y = 0; y < img->height; ++y) {
    pixel_type_t* line = image_row(img, 0, y);
    for (uint16_t x = 0; x < img->width; ++x) {
      BayerIndex bayer_index = GetBayerIndex(kBayerType, x, y);
      switch (bayer_index) {
        case (kR): {
          sum_of_reds += line[x];
          num_of_reds++;
        }; break;
        case (kGr): {
          sum_of_greens += line[x];
          num_of_greens++;
        }; break;
        case (kGb): {
          sum_of_greens += line[x];
          num_of_greens++;
        }; break;
        case (kB): {
          sum_of_blues += line[x];
          num_of_blues++;
        }; break;
        default: {
          break;
        }
      }
    }
  }

  // scale values to green channel
  float average_red = 1.0 * sum_of_reds / num_of_reds;
  float average_green = 1.0 * sum_of_greens / num_of_greens;
  float average_blue = 1.0 * sum_of_blues / num_of_blues;

  float max_average = MAX(MAX(average_red, average_green), average_blue);

  // Convert the float value to fixed point representation, i.e. 0xFF.FF
  uint32_t red_wb = FloatToFixedPoint(max_average / average_red,
                                      kRawPipelineInteger, kRawPipelineFraction,
                                      /*bool is_signed*/ false);
  uint32_t green_wb = FloatToFixedPoint(
      max_average / average_green, kRawPipelineInteger, kRawPipelineFraction,
      /*bool is_signed*/ false);
  uint32_t blue_wb = FloatToFixedPoint(
      max_average / average_blue, kRawPipelineInteger, kRawPipelineFraction,
      /*bool is_signed*/ false);

  wbg_params.gains[0] = red_wb;
  wbg_params.gains[1] = green_wb;
  wbg_params.gains[2] = green_wb;
  wbg_params.gains[3] = blue_wb;
}

void wbg_process(Image* img) {
  if (!wbg_params.enable) return;
  if (!wbg_params.fixed) {
    compute_wbg_gain(img);
  }

  for (uint16_t y = 0; y < img->height; ++y) {
    pixel_type_t* line = image_row(img, 0, y);

    for (uint16_t x = 0; x < img->width; ++x) {
      BayerIndex bayer_index = GetBayerIndex(kBayerType, x, y);
      uint32_t input_val = (uint32_t)line[x];
      uint32_t scaled_pixel = (input_val * wbg_params.gains[bayer_index] +
                               (1 << (kWbgFractional - 1))) >>
                              kWbgFractional;
      line[x] = (pixel_type_t)Clamp(scaled_pixel, kRawPipelineMinVal,
                                    kRawPipelineMaxVal);
    }
  }
}