libheif/nclx.cc

/*
 * HEIF codec.
 * Copyright (c) 2020 Dirk Farin <dirk.farin@gmail.com>
 *
 * This file is part of libheif.
 *
 * libheif is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 *
 * libheif is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with libheif.  If not, see <http://www.gnu.org/licenses/>.
 */


#include "nclx.h"
#include "security_limits.h"

#include <cassert>
#include <limits>
#include <memory>
#include <string>
#include <vector>


primaries::primaries(float gx, float gy, float bx, float by, float rx, float ry, float wx, float wy)
{
  defined = true;
  redX = rx;
  redY = ry;
  greenX = gx;
  greenY = gy;
  blueX = bx;
  blueY = by;
  whiteX = wx;
  whiteY = wy;
}


primaries get_colour_primaries(uint16_t primaries_idx)
{
  switch (primaries_idx) {
    case 1:
      return {0.300f, 0.600f, 0.150f, 0.060f, 0.640f, 0.330f, 0.3127f, 0.3290f};
    case 4:
      return {0.21f, 0.71f, 0.14f, 0.08f, 0.67f, 0.33f, 0.310f, 0.316f};
    case 5:
      return {0.29f, 0.60f, 0.15f, 0.06f, 0.64f, 0.33f, 0.3127f, 0.3290f};
    case 6:
    case 7:
      return {0.310f, 0.595f, 0.155f, 0.070f, 0.630f, 0.340f, 0.3127f, 0.3290f};
    case 8:
      return {0.243f, 0.692f, 0.145f, 0.049f, 0.681f, 0.319f, 0.310f, 0.316f};
    case 9:
      return {0.170f, 0.797f, 0.131f, 0.046f, 0.708f, 0.292f, 0.3127f, 0.3290f};
    case 10:
      return {0.0f, 1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.333333f, 0.33333f};
    case 11:
      return {0.265f, 0.690f, 0.150f, 0.060f, 0.680f, 0.320f, 0.314f, 0.351f};
    case 12:
      return {0.265f, 0.690f, 0.150f, 0.060f, 0.680f, 0.320f, 0.3127f, 0.3290f};
    case 22:
      return {0.295f, 0.605f, 0.155f, 0.077f, 0.630f, 0.340f, 0.3127f, 0.3290f};
    default:
      return {};
  }
}

Kr_Kb Kr_Kb::defaults()
{
  Kr_Kb kr_kb;
  // Rec 601.
  kr_kb.Kr = 0.2990f;
  kr_kb.Kb = 0.1140f;
  return kr_kb;
}


Kr_Kb get_Kr_Kb(uint16_t matrix_coefficients_idx, uint16_t primaries_idx)
{
  Kr_Kb result;

  if (matrix_coefficients_idx == 12 ||
      matrix_coefficients_idx == 13) {

    primaries p = get_colour_primaries(primaries_idx);
    float zr = 1 - (p.redX + p.redY);
    float zg = 1 - (p.greenX + p.greenY);
    float zb = 1 - (p.blueX + p.blueY);
    float zw = 1 - (p.whiteX + p.whiteY);

    float denom = p.whiteY * (p.redX * (p.greenY * zb - p.blueY * zg) + p.greenX * (p.blueY * zr - p.redY * zb) +
                              p.blueX * (p.redY * zg - p.greenY * zr));

    if (denom == 0.0f) {
      return result;
    }

    result.Kr = (p.redY * (p.whiteX * (p.greenY * zb - p.blueY * zg) + p.whiteY * (p.blueX * zg - p.greenX * zb) +
                           zw * (p.greenX * p.blueY - p.blueX * p.greenY))) / denom;
    result.Kb = (p.blueY * (p.whiteX * (p.redY * zg - p.greenY * zr) + p.whiteY * (p.greenX * zr - p.redX * zg) +
                            zw * (p.redX * p.greenY - p.greenX * p.redY))) / denom;
  }
  else
    switch (matrix_coefficients_idx) {
      case 1:
        result.Kr = 0.2126f;
        result.Kb = 0.0722f;
        break;
      case 4:
        result.Kr = 0.30f;
        result.Kb = 0.11f;
        break;
      case 5:
      case 6:
        result.Kr = 0.299f;
        result.Kb = 0.114f;
        break;
      case 7:
        result.Kr = 0.212f;
        result.Kb = 0.087f;
        break;
      case 9:
      case 10:
        result.Kr = 0.2627f;
        result.Kb = 0.0593f;
        break;
      default:;
    }

  return result;
}


YCbCr_to_RGB_coefficients YCbCr_to_RGB_coefficients::defaults()
{
  YCbCr_to_RGB_coefficients coeffs;
  coeffs.defined = true;
  coeffs.r_cr = 1.402f;
  coeffs.g_cb = -0.344136f;
  coeffs.g_cr = -0.714136f;
  coeffs.b_cb = 1.772f;
  return coeffs;
}

YCbCr_to_RGB_coefficients
get_YCbCr_to_RGB_coefficients(uint16_t matrix_coefficients_idx, uint16_t primaries_idx)
{
  YCbCr_to_RGB_coefficients coeffs;

  Kr_Kb k = get_Kr_Kb(matrix_coefficients_idx, primaries_idx);

  if (k.Kb != 0 || k.Kr != 0) { // both will be != 0 when valid
    coeffs.defined = true;
    coeffs.r_cr = 2 * (-k.Kr + 1);
    coeffs.g_cb = 2 * k.Kb * (-k.Kb + 1) / (k.Kb + k.Kr - 1);
    coeffs.g_cr = 2 * k.Kr * (-k.Kr + 1) / (k.Kb + k.Kr - 1);
    coeffs.b_cb = 2 * (-k.Kb + 1);
  }
  else {
    coeffs = YCbCr_to_RGB_coefficients::defaults();
  }

  return coeffs;
}


RGB_to_YCbCr_coefficients
get_RGB_to_YCbCr_coefficients(uint16_t matrix_coefficients_idx, uint16_t primaries_idx)
{
  RGB_to_YCbCr_coefficients coeffs;

  Kr_Kb k = get_Kr_Kb(matrix_coefficients_idx, primaries_idx);

  if (k.Kb != 0 || k.Kr != 0) { // both will be != 0 when valid
    coeffs.defined = true;
    coeffs.c[0][0] = k.Kr;
    coeffs.c[0][1] = 1 - k.Kr - k.Kb;
    coeffs.c[0][2] = k.Kb;
    coeffs.c[1][0] = -k.Kr / (1 - k.Kb) / 2;
    coeffs.c[1][1] = -(1 - k.Kr - k.Kb) / (1 - k.Kb) / 2;
    coeffs.c[1][2] = 0.5f;
    coeffs.c[2][0] = 0.5f;
    coeffs.c[2][1] = -(1 - k.Kr - k.Kb) / (1 - k.Kr) / 2;
    coeffs.c[2][2] = -k.Kb / (1 - k.Kr) / 2;
  }
  else {
    coeffs = RGB_to_YCbCr_coefficients::defaults();
  }

  return coeffs;
}


RGB_to_YCbCr_coefficients RGB_to_YCbCr_coefficients::defaults()
{
  RGB_to_YCbCr_coefficients coeffs;
  coeffs.defined = true;
  // Rec 601 full. Kr=0.2990f Kb=0.1140f.
  coeffs.c[0][0] = 0.299f;
  coeffs.c[0][1] = 0.587f;
  coeffs.c[0][2] = 0.114f;
  coeffs.c[1][0] = -0.168735f;
  coeffs.c[1][1] = -0.331264f;
  coeffs.c[1][2] = 0.5f;
  coeffs.c[2][0] = 0.5f;
  coeffs.c[2][1] = -0.418688f;
  coeffs.c[2][2] = -0.081312f;

  return coeffs;
}


Error color_profile_nclx::parse(BitstreamRange& range)
{
  StreamReader::grow_status status;
  status = range.wait_for_available_bytes(7);
  if (status != StreamReader::size_reached) {
    // TODO: return recoverable error at timeout
    return Error(heif_error_Invalid_input,
                 heif_suberror_End_of_data);
  }

  m_colour_primaries = range.read16();
  m_transfer_characteristics = range.read16();
  m_matrix_coefficients = range.read16();
  m_full_range_flag = (range.read8() & 0x80 ? true : false);

  return Error::Ok;
}

Error color_profile_nclx::get_nclx_color_profile(struct heif_color_profile_nclx** out_data) const
{
  *out_data = nullptr;

  struct heif_color_profile_nclx* nclx = alloc_nclx_color_profile();

  if (nclx == nullptr) {
    return Error(heif_error_Memory_allocation_error,
                 heif_suberror_Unspecified);
  }

  struct heif_error err;

  nclx->version = 1;

  err = heif_nclx_color_profile_set_color_primaries(nclx, get_colour_primaries());
  if (err.code) {
    free_nclx_color_profile(nclx);
    return {err.code, err.subcode};
  }

  err = heif_nclx_color_profile_set_transfer_characteristics(nclx, get_transfer_characteristics());
  if (err.code) {
    free_nclx_color_profile(nclx);
    return {err.code, err.subcode};
  }

  err = heif_nclx_color_profile_set_matrix_coefficients(nclx, get_matrix_coefficients());
  if (err.code) {
    free_nclx_color_profile(nclx);
    return {err.code, err.subcode};
  }

  nclx->full_range_flag = get_full_range_flag();

  // fill color primaries

  auto primaries = ::get_colour_primaries(nclx->color_primaries);

  nclx->color_primary_red_x = primaries.redX;
  nclx->color_primary_red_y = primaries.redY;
  nclx->color_primary_green_x = primaries.greenX;
  nclx->color_primary_green_y = primaries.greenY;
  nclx->color_primary_blue_x = primaries.blueX;
  nclx->color_primary_blue_y = primaries.blueY;
  nclx->color_primary_white_x = primaries.whiteX;
  nclx->color_primary_white_y = primaries.whiteY;

  *out_data = nclx;

  return Error::Ok;
}


struct heif_color_profile_nclx* color_profile_nclx::alloc_nclx_color_profile()
{
  auto profile = (heif_color_profile_nclx*) malloc(sizeof(struct heif_color_profile_nclx));

  if (profile) {
    profile->version = 1;

    // sRGB defaults
    profile->color_primaries = heif_color_primaries_ITU_R_BT_709_5; // 1
    profile->transfer_characteristics = heif_transfer_characteristic_IEC_61966_2_1; // 13
    profile->matrix_coefficients = heif_matrix_coefficients_ITU_R_BT_601_6; // 6
    profile->full_range_flag = true;
  }

  return profile;
}


void color_profile_nclx::free_nclx_color_profile(struct heif_color_profile_nclx* profile)
{
  free(profile);
}


void color_profile_nclx::set_sRGB_defaults()
{
  // sRGB defaults
  m_colour_primaries = 1;
  m_transfer_characteristics = 13;
  m_matrix_coefficients = 6;
  m_full_range_flag = true;
}


void color_profile_nclx::set_undefined()
{
  m_colour_primaries = 2;
  m_transfer_characteristics = 2;
  m_matrix_coefficients = 2;
  m_full_range_flag = true;
}


void color_profile_nclx::set_from_heif_color_profile_nclx(const struct heif_color_profile_nclx* nclx)
{
  if (nclx) {
    m_colour_primaries = nclx->color_primaries;
    m_transfer_characteristics = nclx->transfer_characteristics;
    m_matrix_coefficients = nclx->matrix_coefficients;
    m_full_range_flag = nclx->full_range_flag;
  }
}


void color_profile_nclx::replace_undefined_values_with_sRGB_defaults()
{
  if (m_matrix_coefficients == heif_matrix_coefficients_unspecified) {
    m_matrix_coefficients = heif_matrix_coefficients_ITU_R_BT_601_6;
  }

  if (m_colour_primaries == heif_color_primaries_unspecified) {
    m_colour_primaries = heif_color_primaries_ITU_R_BT_709_5;
  }

  if (m_transfer_characteristics == heif_transfer_characteristic_unspecified) {
    m_transfer_characteristics = heif_transfer_characteristic_IEC_61966_2_1;
  }
}


Error Box_colr::parse(BitstreamRange& range)
{
  StreamReader::grow_status status;
  uint32_t colour_type = range.read32();

  if (colour_type == fourcc("nclx")) {
    auto color_profile = std::make_shared<color_profile_nclx>();
    m_color_profile = color_profile;
    Error err = color_profile->parse(range);
    if (err) {
      return err;
    }
  }
  else if (colour_type == fourcc("prof") ||
           colour_type == fourcc("rICC")) {
    if (!has_fixed_box_size()) {
      return Error(heif_error_Unsupported_feature, heif_suberror_Unspecified, "colr boxes with undefined box size are not supported");
    }

    uint64_t profile_size_64 = get_box_size() - get_header_size() - 4;
    if (profile_size_64 > MAX_COLOR_PROFILE_SIZE) {
      return Error(heif_error_Invalid_input, heif_suberror_Security_limit_exceeded, "Color profile exceeds maximum supported size");
    }

    size_t profile_size = static_cast<size_t>(profile_size_64);

    status = range.wait_for_available_bytes(profile_size);
    if (status != StreamReader::size_reached) {
      // TODO: return recoverable error at timeout
      return Error(heif_error_Invalid_input,
                   heif_suberror_End_of_data);
    }

    std::vector<uint8_t> rawData(profile_size);
    for (size_t i = 0; i < profile_size; i++) {
      rawData[i] = range.read8();
    }

    m_color_profile = std::make_shared<color_profile_raw>(colour_type, rawData);
  }
  else {
    return Error(heif_error_Invalid_input,
                 heif_suberror_Unknown_color_profile_type);
  }

  return range.get_error();
}


std::string Box_colr::dump(Indent& indent) const
{
  std::ostringstream sstr;
  sstr << Box::dump(indent);

  if (m_color_profile) {
    sstr << indent << "colour_type: " << to_fourcc(get_color_profile_type()) << "\n";
    sstr << m_color_profile->dump(indent);
  }
  else {
    sstr << indent << "colour_type: ---\n";
    sstr << "no color profile\n";
  }

  return sstr.str();
}


std::string color_profile_raw::dump(Indent& indent) const
{
  std::ostringstream sstr;
  sstr << indent << "profile size: " << m_data.size() << "\n";
  return sstr.str();
}


std::string color_profile_nclx::dump(Indent& indent) const
{
  std::ostringstream sstr;
  sstr << indent << "colour_primaries: " << m_colour_primaries << "\n"
       << indent << "transfer_characteristics: " << m_transfer_characteristics << "\n"
       << indent << "matrix_coefficients: " << m_matrix_coefficients << "\n"
       << indent << "full_range_flag: " << m_full_range_flag << "\n";
  return sstr.str();
}


Error color_profile_nclx::write(StreamWriter& writer) const
{
  writer.write16(m_colour_primaries);
  writer.write16(m_transfer_characteristics);
  writer.write16(m_matrix_coefficients);
  writer.write8(m_full_range_flag ? 0x80 : 0x00);

  return Error::Ok;
}

Error color_profile_raw::write(StreamWriter& writer) const
{
  writer.write(m_data);

  return Error::Ok;
}

Error Box_colr::write(StreamWriter& writer) const
{
  size_t box_start = reserve_box_header_space(writer);

  assert(m_color_profile);

  writer.write32(m_color_profile->get_type());

  Error err = m_color_profile->write(writer);
  if (err) {
    return err;
  }

  prepend_header(writer, box_start);

  return Error::Ok;
}