blob: f38135a6738da093615a590c9075c168fce30646 [file] [log] [blame]
#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* 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: {
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* 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);
}
}
}