risp4ml/isp_stages/downscale.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/downscale.h"

#include "risp4ml/common/utils.h"

static const uint16_t kScalePrecision = 8;
static const uint16_t kScaleFixedOne = (1 << kScalePrecision);

static DownscaleParams params = {
    .enable = true,
    .scale_precision = kScalePrecision,
    .interpolate_precision = 8,
    .interpolate_shift = 2,
    .scale_fixed_one = kScaleFixedOne,
    .scale_fraction_mask = kScaleFixedOne - 1,
    .weight_shift = 0,
    .hor_scale_factor = kScaleFixedOne,
    .ver_scale_factor = kScaleFixedOne,
    .ver_initial_offset = 0,
    .hor_initial_offset = 0,
};

void set_downscale_param(DownscaleParams* in_params) { params = *in_params; }

void set_downscale_factor(Image* input, ImageU8* output) {
  params.hor_scale_factor =
      (kScaleFixedOne * input->width - params.hor_initial_offset) /
      output->width;
  params.ver_scale_factor =
      (kScaleFixedOne * input->height - params.ver_initial_offset) /
      output->height;
}

// Basic bilinear downscale

// Resamples image using bilinear interpolation.
// 'output' is modified by this function to store the output image.
void downscale_process(Image* input, ImageU8* output) {
  if (!params.enable) {
    return;
  }
  for (uint16_t channel = 0; channel < output->num_channels; ++channel) {
    // Each output pixel at (x, y) is sampled at (X, Y) in the input image
    // coordinate according to the following formula:
    //    Y = (y * ver_scale_factor) + ver_initial_offset
    //    X = (x * hor_scale_factor) + hor_initial_offset
    // `accumulated_pos_*` is the location in the inout image calculated by
    // repeated addition of *_scale_factor
    // `integer_pos_*` is the nearest integer index at top/left side.
    uint32_t accumulated_pos_y = params.ver_initial_offset;
    for (uint16_t y = 0; y < output->height; ++y) {
      uint32_t integer_pos_y = accumulated_pos_y >> params.scale_precision;
      uint32_t y0 = Clamp(integer_pos_y, 0, input->height - 1);
      uint32_t y1 = Clamp(integer_pos_y + 1, 0, input->height - 1);
      // The fractional part of the accumulated position gives us the weight but
      // in scale_precision, since we want the weights to be in
      // interpolate_precision we shift by their bitwidth difference which is
      // weight_shift
      uint32_t weight_y = Round(accumulated_pos_y & params.scale_fraction_mask,
                                params.weight_shift);
      uint32_t accumulated_pos_x = params.hor_initial_offset;

      for (uint16_t x = 0; x < output->width; ++x) {
        uint32_t integer_pos_x = accumulated_pos_x >> params.scale_precision;
        uint32_t x0 = Clamp(integer_pos_x, 0, input->width - 1);
        uint32_t x1 = Clamp(integer_pos_x + 1, 0, input->width - 1);
        uint32_t weight_x =
            Round(accumulated_pos_x & params.scale_fraction_mask,
                  params.weight_shift);

        // Perform vertical interpolation first to get p0 and p1,
        // then horizontal interpolation to get output value.
        // `interpolate_shift` is for preserving floating point for interpolated
        // values to avoid incremented quantization errors.
        //   (x0, y0)             (x1, y0)
        //      |                    |
        //      v                    v
        //      p0  ->  output  <-   p1
        //      ^                    ^
        //      |                    |
        //   (x0, y1)             (x1, y1)

        uint32_t p0 = Lerp(
            image_pixel_val(input, channel, y0, x0) << params.interpolate_shift,
            image_pixel_val(input, channel, y1, x0) << params.interpolate_shift,
            weight_y, params.interpolate_precision);
        uint32_t p1 = Lerp(
            image_pixel_val(input, channel, y0, x1) << params.interpolate_shift,
            image_pixel_val(input, channel, y1, x1) << params.interpolate_shift,
            weight_y, params.interpolate_precision);
        // To normalize the output we shift right back by interpolate shift
        uint32_t tmp_interpolated =
            Round(Lerp(p0, p1, weight_x, params.interpolate_precision),
                  params.interpolate_shift);

        // As this is the final stage and we'd like to have the output image in
        // an interleaved RGBRGBRGB format.
        // The accessor (*output)(channel, y, x)
        // is assuming planar layout of RGB channels - uses stride_c_ which is
        // set in the constructor to stride_c_ = width * height.
        // Hence image data is accessed and manipulated directly.
        // Also should avoid using stride_c, stride_x and stride_y as they are
        // defined in the image class with the planar layout assumption.
        const uint16_t interleaved_stride_c = 1;
        const uint16_t interleaved_stride_x = output->num_channels;
        const uint32_t interleaved_stride_y =
            output->num_channels * output->width;

        // Shift interpolated result to output bitwidth. Not rounding to avoid
        // overflow to 256.
        output->data[y * interleaved_stride_y + x * interleaved_stride_x +
                     channel * interleaved_stride_c] =
            (uint8_t)(tmp_interpolated >>
                      (kRawPipelineBpp - kPipeOutputBpp));  // 16 - 8

        accumulated_pos_x += params.hor_scale_factor;
      }
      accumulated_pos_y += params.ver_scale_factor;
    }
  }
}