samples/risp4ml/isp_stages/wbg.c

#include <assert.h>

#include "samples/risp4ml/common/utils.h"
#include "samples/risp4ml/isp_stages/wbg.h"

#define kBayerColorChannels 4
#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* input) {
  // Calculate the white-balance gain values using the "gray world" algorithm
  uint16_t height = input->height;
  uint16_t width = input->width;

  pixel_type_t* in_line;

  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 < height; ++y) {
    in_line = image_row(input, 0, y);
    for (uint16_t x = 0; x < width; ++x) {
      BayerIndex bayer_index = GetBayerIndex(kBayerType, x, y);
      switch (bayer_index) {
        case (kR): {
          sum_of_reds += in_line[x];
          num_of_reds++;
        }; break;
        case (kGr): {
          sum_of_greens += in_line[x];
          num_of_greens++;
        }; break;
        case (kGb): {
          sum_of_greens += in_line[x];
          num_of_greens++;
        }; break;
        case (kB): {
          sum_of_blues += in_line[x];
          num_of_blues++;
        }; break;
        default: {
          assert(0 && "Unexpected channel index");
        }
      }
    }
  }

  // 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* input, Image* output) {
  if (!wbg_params.enable) {
    *output = *input;
    return;
  }

  uint16_t height = input->height;
  uint16_t width = input->width;

  const pixel_type_t* in_line;
  pixel_type_t* out_line;

  if (!wbg_params.fixed) {
    compute_wbg_gain(input);
  }

  for (uint16_t y = 0; y < height; ++y) {
    in_line = image_row(input, 0, y);
    out_line = image_row(output, 0, y);

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