samples/risp4ml/isp_stages/demosaic.c

/*
 * Copyright 2022 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 <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");
        }
      }
    }
  }
}