samples/risp4ml/isp_stages/demosaic.c

#include <assert.h>

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

#define kRgbColorChannels 3

static DemosaicParams demosaic_params = {.enable = true};

void set_demosaic_params(DemosaicParams* params) { demosaic_params = *params; }

// Basic bilinear demosaic
void demosaic_process(Image* input, Image* output) {
  if (!demosaic_params.enable) {
    return;
  }
  uint16_t height = input->height;
  uint16_t width = input->width;

  const pixel_type_t* line_buffers[kRgbColorChannels];
  int x_offset[kRgbColorChannels];

  for (uint16_t y = 0; y < height; ++y) {
    line_buffers[0] = (y) ? image_row(input, 0, y - 1) : image_row(input, 0, 1);
    line_buffers[1] = image_row(input, 0, y);
    line_buffers[2] = (y < height - 1) ? image_row(input, 0, y + 1)
                                       : image_row(input, 0, height - 2);

    for (uint16_t x = 0; x < width; ++x) {
      for (uint16_t c = 0; c < kRgbColorChannels; ++c) {
        x_offset[c] = BayerMirrorBoundary(x - 1 + c, width);
      }

      BayerIndex bayer_index = GetBayerIndex(kBayerType, x, y);
      switch (bayer_index) {
        case (kR): {
          *image_pixel(output, 0, y, x) = line_buffers[1][x_offset[1]];
          *image_pixel(output, 1, y, x) =
              (line_buffers[0][x_offset[1]] + line_buffers[2][x_offset[1]] +
               line_buffers[1][x_offset[0]] + line_buffers[1][x_offset[2]]) /
              4;
          *image_pixel(output, 2, y, x) =
              (line_buffers[0][x_offset[0]] + line_buffers[0][x_offset[2]] +
               line_buffers[2][x_offset[0]] + line_buffers[2][x_offset[2]]) /
              4;
        }; break;
        case (kGr): {
          *image_pixel(output, 0, y, x) =
              (line_buffers[1][x_offset[0]] + line_buffers[1][x_offset[2]]) / 2;
          *image_pixel(output, 1, y, x) = line_buffers[1][x_offset[1]];
          *image_pixel(output, 2, y, x) =
              (line_buffers[0][x_offset[1]] + line_buffers[2][x_offset[1]]) / 2;
        }; break;
        case (kGb): {
          *image_pixel(output, 0, y, x) =
              (line_buffers[0][x_offset[1]] + line_buffers[2][x_offset[1]]) / 2;
          *image_pixel(output, 1, y, x) = line_buffers[1][x_offset[1]];
          *image_pixel(output, 2, y, x) =
              (line_buffers[1][x_offset[0]] + line_buffers[1][x_offset[2]]) / 2;
        }; break;
        case (kB): {
          *image_pixel(output, 0, y, x) =
              (line_buffers[0][x_offset[0]] + line_buffers[0][x_offset[2]] +
               line_buffers[2][x_offset[0]] + line_buffers[2][x_offset[2]]) /
              4;
          *image_pixel(output, 1, y, x) =
              (line_buffers[0][x_offset[1]] + line_buffers[2][x_offset[1]] +
               line_buffers[1][x_offset[0]] + line_buffers[1][x_offset[2]]) /
              4;
          *image_pixel(output, 2, y, x) = line_buffers[1][x_offset[1]];
        }; break;
        default: {
          assert(0 && "Unexpected channel index");
        }
      }
    }
  }
}