quantization.cpp

// Aseprite Render Library
// Copyright (c) 2019-2022  Igara Studio S.A.
// Copyright (c) 2001-2018  David Capello
//
// This file is released under the terms of the MIT license.
// Read LICENSE.txt for more information.

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "render/quantization.h"

#include "doc/image_impl.h"
#include "doc/layer.h"
#include "doc/octree_map.h"
#include "doc/palette.h"
#include "doc/primitives.h"
#include "doc/remap.h"
#include "doc/sprite.h"
#include "render/dithering.h"
#include "render/error_diffusion.h"
#include "render/ordered_dither.h"
#include "render/render.h"
#include "render/task_delegate.h"

#include <algorithm>
#include <limits>
#include <map>
#include <memory>
#include <vector>

namespace render {

using namespace doc;
using namespace gfx;

Palette* create_palette_from_sprite(
  const Sprite* sprite,
  const frame_t fromFrame,
  const frame_t toFrame,
  const bool withAlpha,
  Palette* palette,
  TaskDelegate* delegate,
  const bool newBlend,
  RgbMapAlgorithm mapAlgo,
  const bool calculateWithTransparent)
{
   if (mapAlgo == doc::RgbMapAlgorithm::DEFAULT)
     mapAlgo = doc::RgbMapAlgorithm::OCTREE;

  PaletteOptimizer optimizer;
  OctreeMap octreemap;

  // Transparent color is needed if we have transparent layers
  int maskIndex;
  if ((sprite->backgroundLayer() && sprite->allLayersCount() == 1) ||
      !calculateWithTransparent)
    maskIndex = -1;
  else if (sprite->colorMode() == ColorMode::INDEXED)
    maskIndex = sprite->transparentColor();
  else {
    ASSERT(sprite->transparentColor() == 0);
    maskIndex = 0; // For RGB/Grayscale images we use index 0 as the transparent index by default
  }

  // TODO check if how this is used in OctreeMap, if as a RGBA value
  //      or as an index (here we are using it as an index).
  const color_t maskColor = (sprite->backgroundLayer()
                             && sprite->allLayersCount() == 1) ? DOC_OCTREE_IS_OPAQUE:
                                                                 sprite->transparentColor();

  if (!palette)
    palette = new Palette(fromFrame, 256);

  // Add a flat image with the current sprite's frame rendered
  ImageRef flat_image(Image::create(IMAGE_RGB,
      sprite->width(), sprite->height()));

  render::Render render;
  render.setNewBlend(newBlend);

  // Feed the optimizer with all rendered frames
  for (frame_t frame=fromFrame; frame<=toFrame; ++frame) {
    render.renderSprite(flat_image.get(), sprite, frame);

    switch (mapAlgo) {
      case RgbMapAlgorithm::RGB5A3:
        optimizer.feedWithImage(flat_image.get(), withAlpha);
        break;
      case RgbMapAlgorithm::OCTREE:
        octreemap.feedWithImage(flat_image.get(), withAlpha, maskColor);
        break;
      default:
        ASSERT(false);
        break;
    }

    if (delegate) {
      if (!delegate->continueTask())
        return nullptr;

      delegate->notifyTaskProgress(
        double(frame-fromFrame+1) / double(toFrame-fromFrame+1));
    }
  }

  switch (mapAlgo) {

    case RgbMapAlgorithm::RGB5A3: {
      // Generate an optimized palette
      optimizer.calculate(palette, maskIndex);
      break;
    }

    case RgbMapAlgorithm::OCTREE:
      // TODO check calculateWithTransparent flag

      if (!octreemap.makePalette(palette, palette->size())) {
        // We can use an 8-bit deep octree map, instead of 7-bit of the
        // first attempt.
        octreemap = OctreeMap();
        for (frame_t frame=fromFrame; frame<=toFrame; ++frame) {
          render.renderSprite(flat_image.get(), sprite, frame);
          octreemap.feedWithImage(flat_image.get(), withAlpha, maskColor , 8);
          if (delegate) {
            if (!delegate->continueTask())
              return nullptr;

            delegate->notifyTaskProgress(
              double(frame-fromFrame+1) / double(toFrame-fromFrame+1));
          }
        }
        octreemap.makePalette(palette, palette->size(), 8);
      }
      break;
  }

  return palette;
}

Image* convert_pixel_format(
  const Image* image,
  Image* new_image,
  PixelFormat pixelFormat,
  const Dithering& dithering,
  const RgbMap* rgbmap,
  const Palette* palette,
  bool is_background,
  color_t new_mask_color,
  rgba_to_graya_func toGray,
  TaskDelegate* delegate)
{
  if (!new_image)
    new_image = Image::create(pixelFormat, image->width(), image->height());
  new_image->setMaskColor(new_mask_color);

  // RGB -> Indexed with ordered dithering
  if (image->pixelFormat() == IMAGE_RGB &&
      pixelFormat == IMAGE_INDEXED &&
      dithering.algorithm() != DitheringAlgorithm::None) {
    std::unique_ptr<DitheringAlgorithmBase> dither;
    switch (dithering.algorithm()) {
      case DitheringAlgorithm::Ordered:
        dither.reset(new OrderedDither2(is_background ? -1: new_mask_color));
        break;
      case DitheringAlgorithm::Old:
        dither.reset(new OrderedDither(is_background ? -1: new_mask_color));
        break;
      case DitheringAlgorithm::ErrorDiffusion:
        dither.reset(new ErrorDiffusionDither(is_background ? -1: new_mask_color));
        break;
    }
    if (dither)
      dither_rgb_image_to_indexed(
        *dither, dithering,
        image, new_image, rgbmap, palette, delegate);
    return new_image;
  }

  // RGB/Indexed -> Gray
  if ((image->pixelFormat() == IMAGE_RGB ||
       image->pixelFormat() == IMAGE_INDEXED) &&
      new_image->pixelFormat() == IMAGE_GRAYSCALE) {
    if (!toGray)
      toGray = &rgba_to_graya_using_luma;
  }

  color_t c;
  int r, g, b, a;

  switch (image->pixelFormat()) {

    case IMAGE_RGB: {
      const LockImageBits<RgbTraits> srcBits(image);
      auto src_it = srcBits.begin(), src_end = srcBits.end();

      switch (new_image->pixelFormat()) {

        // RGB -> RGB
        case IMAGE_RGB:
          new_image->copy(image, gfx::Clip(image->bounds()));
          break;

        // RGB -> Grayscale
        case IMAGE_GRAYSCALE: {
          LockImageBits<GrayscaleTraits> dstBits(new_image, Image::WriteLock);
          auto dst_it = dstBits.begin();
#ifdef _DEBUG
          auto dst_end = dstBits.end();
          ASSERT(toGray);
#endif
          for (; src_it != src_end; ++src_it, ++dst_it) {
            ASSERT(dst_it != dst_end);
            *dst_it = (*toGray)(*src_it);
          }
          ASSERT(dst_it == dst_end);
          break;
        }

        // RGB -> Indexed
        case IMAGE_INDEXED: {
          LockImageBits<IndexedTraits> dstBits(new_image, Image::WriteLock);
          auto dst_it = dstBits.begin();
#ifdef _DEBUG
          auto dst_end = dstBits.end();
#endif

          for (; src_it != src_end; ++src_it, ++dst_it) {
            ASSERT(dst_it != dst_end);
            c = *src_it;

            r = rgba_getr(c);
            g = rgba_getg(c);
            b = rgba_getb(c);
            a = rgba_geta(c);

            if (a == 0)
              *dst_it = (new_mask_color == -1? 0 : new_mask_color);
            else if (rgbmap)
              *dst_it = rgbmap->mapColor(c);
            else
              *dst_it = palette->findBestfit(r, g, b, a, new_mask_color);
          }
          ASSERT(dst_it == dst_end);
          break;
        }
      }
      break;
    }

    case IMAGE_GRAYSCALE: {
      const LockImageBits<GrayscaleTraits> srcBits(image);
      auto src_it = srcBits.begin(), src_end = srcBits.end();

      switch (new_image->pixelFormat()) {

        // Grayscale -> RGB
        case IMAGE_RGB: {
          LockImageBits<RgbTraits> dstBits(new_image, Image::WriteLock);
          auto dst_it = dstBits.begin();
#ifdef _DEBUG
          auto dst_end = dstBits.end();
#endif
          for (; src_it != src_end; ++src_it, ++dst_it) {
            ASSERT(dst_it != dst_end);
            c = *src_it;

            g = graya_getv(c);

            *dst_it = rgba(g, g, g, graya_geta(c));
          }
          ASSERT(dst_it == dst_end);
          break;
        }

        // Grayscale -> Grayscale
        case IMAGE_GRAYSCALE:
          new_image->copy(image, gfx::Clip(image->bounds()));
          break;

        // Grayscale -> Indexed
        case IMAGE_INDEXED: {
          LockImageBits<IndexedTraits> dstBits(new_image, Image::WriteLock);
          auto dst_it = dstBits.begin();
#ifdef _DEBUG
          auto dst_end = dstBits.end();
#endif
          for (; src_it != src_end; ++src_it, ++dst_it) {
            ASSERT(dst_it != dst_end);
            c = *src_it;
            a = graya_geta(c);
            c = graya_getv(c);

            if (a == 0)
              *dst_it = (new_mask_color == -1? 0 : new_mask_color);
            else if (rgbmap)
              *dst_it = rgbmap->mapColor(c, c, c, a);
            else
              *dst_it = palette->findBestfit(c, c, c, a, new_mask_color);
          }
          ASSERT(dst_it == dst_end);
          break;
        }
      }
      break;
    }

    case IMAGE_INDEXED: {
      const LockImageBits<IndexedTraits> srcBits(image);
      auto src_it = srcBits.begin(), src_end = srcBits.end();

      switch (new_image->pixelFormat()) {

        // Indexed -> RGB
        case IMAGE_RGB: {
          LockImageBits<RgbTraits> dstBits(new_image, Image::WriteLock);
          auto dst_it = dstBits.begin();
#ifdef _DEBUG
          auto dst_end = dstBits.end();
#endif
          for (; src_it != src_end; ++src_it, ++dst_it) {
            ASSERT(dst_it != dst_end);
            c = *src_it;

            if (!is_background && c == image->maskColor())
              *dst_it = rgba(0, 0, 0, 0);
            else {
              const uint32_t p = palette->getEntry(c);
              if (is_background)
                *dst_it = rgba(rgba_getr(p), rgba_getg(p), rgba_getb(p), 255);
              else
                *dst_it = p;
            }
          }
          ASSERT(dst_it == dst_end);
          break;
        }

        // Indexed -> Grayscale
        case IMAGE_GRAYSCALE: {
          LockImageBits<GrayscaleTraits> dstBits(new_image, Image::WriteLock);
          auto dst_it = dstBits.begin();
#ifdef _DEBUG
          auto dst_end = dstBits.end();
          ASSERT(toGray);
#endif
          for (; src_it != src_end; ++src_it, ++dst_it) {
            ASSERT(dst_it != dst_end);
            c = *src_it;

            if (!is_background && c == image->maskColor())
              *dst_it = graya(0, 0);
            else {
              c = palette->getEntry(c);
              *dst_it = (*toGray)(c);
            }
          }
          ASSERT(dst_it == dst_end);
          break;
        }

        // Indexed -> Indexed
        case IMAGE_INDEXED: {
          LockImageBits<IndexedTraits> dstBits(new_image, Image::WriteLock);
          auto dst_it = dstBits.begin();
#ifdef _DEBUG
          auto dst_end = dstBits.end();
#endif
          for (; src_it != src_end; ++src_it, ++dst_it) {
            ASSERT(dst_it != dst_end);
            c = *src_it;

            if (!is_background && c == image->maskColor())
              *dst_it = new_mask_color;
            else {
              c = palette->getEntry(c);
              r = rgba_getr(c);
              g = rgba_getg(c);
              b = rgba_getb(c);
              a = rgba_geta(c);

              if (rgbmap)
                *dst_it = rgbmap->mapColor(r, g, b, a);
              else
                *dst_it = palette->findBestfit(r, g, b, a, new_mask_color);
            }
          }
          ASSERT(dst_it == dst_end);
          break;
        }

      }
      break;
    }
  }

  return new_image;
}

//////////////////////////////////////////////////////////////////////
// Creation of optimized palette for RGB images
// by David Capello

void PaletteOptimizer::feedWithImage(const Image* image,
                                     const bool withAlpha)
{
  feedWithImage(image, image->bounds(), withAlpha);
}

void PaletteOptimizer::feedWithImage(const Image* image,
                                     const gfx::Rect& bounds,
                                     const bool withAlpha)
{
  uint32_t color;

  if (withAlpha)
    m_withAlpha = true;

  ASSERT(image);
  switch (image->pixelFormat()) {

    case IMAGE_RGB:
      {
        const LockImageBits<RgbTraits> bits(image, bounds);
        auto it = bits.begin(), end = bits.end();

        for (; it != end; ++it) {
          color = *it;
          if (rgba_geta(color) > 0) {
            if (!withAlpha)
              color |= rgba(0, 0, 0, 255);

            m_histogram.addSamples(color, 1);
          }
        }
      }
      break;

    case IMAGE_GRAYSCALE:
      {
        const LockImageBits<GrayscaleTraits> bits(image, bounds);
        auto it = bits.begin(), end = bits.end();

        for (; it != end; ++it) {
          color = *it;

          if (graya_geta(color) > 0) {
            if (!withAlpha)
              color = graya(graya_getv(color), 255);

            m_histogram.addSamples(rgba(graya_getv(color),
                                        graya_getv(color),
                                        graya_getv(color),
                                        graya_geta(color)), 1);
          }
        }
      }
      break;

    case IMAGE_INDEXED:
      ASSERT(false);
      break;

  }
}

void PaletteOptimizer::feedWithRgbaColor(color_t color)
{
  m_histogram.addSamples(color, 1);
}

void PaletteOptimizer::calculate(Palette* palette, int maskIndex)
{
  bool addMask;

  if ((palette->size() > 1) &&
      (maskIndex >= 0 && maskIndex < palette->size())) {
    palette->resize(palette->size()-1);
    addMask = true;
  }
  else
    addMask = false;

  // If the sprite has a background layer, the first entry can be
  // used, in other case the 0 indexed will be the mask color, so it
  // will not be used later in the color conversion (from RGB to
  // Indexed).
  int usedColors = m_histogram.createOptimizedPalette(palette);

  if (addMask) {
    palette->resize(usedColors+1);

    Remap remap(palette->size());
    for (int i=0; i<usedColors; ++i)
      remap.map(i, i + (i >= maskIndex ? 1: 0));

    palette->applyRemap(remap);

    if (maskIndex < palette->size())
      palette->setEntry(maskIndex, rgba(0, 0, 0, (m_withAlpha ? 0: 255)));
  }
  else {
    palette->resize(std::max(1, usedColors));
  }
}

} // namespace render