| """ |
| Copyright (c) 2024, Alliance for Open Media. All rights reserved |
| |
| This source code is subject to the terms of the BSD 3-Clause Clear License |
| and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear |
| License was not distributed with this source code in the LICENSE file, you |
| can obtain it at aomedia.org/license/software-license/bsd-3-c-c/. If the |
| Alliance for Open Media Patent License 1.0 was not distributed with this |
| source code in the PATENTS file, you can obtain it at |
| aomedia.org/license/patent-license/. |
| """ |
| import numpy as np |
| |
| from parakit.entropy.codec_default_cdf import ( |
| AV1_PROB_COST, |
| AV1_PROB_COST_SHIFT, |
| CDF_INIT_TOP, |
| CDF_PROB_BITS, |
| CDF_PROB_TOP, |
| ) |
| |
| |
| def clog2(x): |
| """Ceiling of log2""" |
| if x <= 0: |
| raise ValueError("clog2 input error") |
| return (x - 1).bit_length() |
| |
| |
| def flog2(x): |
| """Floor of log2""" |
| if x <= 0: |
| raise ValueError("flog2 input error") |
| return x.bit_length() - 1 |
| |
| |
| def update_cdfinv_av1(cdf, val, counter, nsymb, roffset=0): |
| """Python implementation of the following C code from AVM codec: |
| -------------------------------------------------------------- |
| static INLINE void update_cdf(aom_cdf_prob *cdf, int8_t val, int nsymbs) { |
| int rate; |
| int i, tmp; |
| static const int nsymbs2speed[17] = { 0, 0, 1, 1, 2, 2, 2, 2, 2, |
| 2, 2, 2, 2, 2, 2, 2, 2 }; |
| assert(nsymbs < 17); |
| rate = 3 + (cdf[nsymbs] > 15) + (cdf[nsymbs] > 31) + |
| nsymbs2speed[nsymbs]; // + get_msb(nsymbs); |
| tmp = AOM_ICDF(0); |
| // Single loop (faster) |
| for (i = 0; i < nsymbs - 1; ++i) { |
| tmp = (i == val) ? 0 : tmp; |
| if (tmp < cdf[i]) { |
| cdf[i] -= ((cdf[i] - tmp) >> rate); |
| } else { |
| cdf[i] += ((tmp - cdf[i]) >> rate); |
| } |
| } |
| cdf[nsymbs] += (cdf[nsymbs] < 32); |
| } |
| -------------------------------------------------------------- |
| """ |
| nsymbs2speed = [0, 0, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2] |
| rate = 3 + nsymbs2speed[nsymb] |
| if counter > 15: |
| rate = rate + 1 |
| if counter > 31: |
| rate = rate + 1 |
| rate = rate + roffset |
| tmp = CDF_INIT_TOP |
| for i in range(nsymb - 1): |
| if i == val: |
| tmp = 0 |
| if tmp < cdf[i]: |
| cdf[i] -= (cdf[i] - tmp) // (2**rate) |
| else: |
| cdf[i] += (tmp - cdf[i]) // (2**rate) |
| return cdf |
| |
| |
| def get_prob(num, den): |
| """Python implementation of the following C code from AVM codec: |
| -------------------------------------------------------------- |
| static INLINE uint8_t get_prob(unsigned int num, unsigned int den) { |
| assert(den != 0); |
| { |
| const int p = (int)(((uint64_t)num * 256 + (den >> 1)) / den); |
| // (p > 255) ? 255 : (p < 1) ? 1 : p; |
| const int clipped_prob = p | ((255 - p) >> 23) | (p == 0); |
| return (uint8_t)clipped_prob; |
| } |
| } |
| -------------------------------------------------------------- |
| """ |
| p = int(((num * 256) + (den // 2)) / den) |
| if p > 255: |
| p = 255 |
| if p < 1: |
| p = 1 |
| return p |
| |
| |
| def cost_literal_av1(n): |
| """Python implementation of the following C code from AVM codec: |
| -------------------------------------------------------------- |
| define av1_cost_literal(n) ((n) * (1 << AV1_PROB_COST_SHIFT)) |
| -------------------------------------------------------------- |
| """ |
| return n * (2**AV1_PROB_COST_SHIFT) |
| |
| |
| def cost_symbol_av1(p15): |
| """Python implementation of the following C code from AVM codec: |
| -------------------------------------------------------------- |
| static INLINE int av1_cost_symbol(aom_cdf_prob p15) { |
| // p15 can be out of range [1, CDF_PROB_TOP - 1]. Clamping it, so that the |
| // following cost calculation works correctly. Otherwise, if p15 = |
| // CDF_PROB_TOP, shift would be -1, and "p15 << shift" would be wrong. |
| p15 = (aom_cdf_prob)clamp(p15, 1, CDF_PROB_TOP - 1); |
| assert(0 < p15 && p15 < CDF_PROB_TOP); |
| const int shift = CDF_PROB_BITS - 1 - get_msb(p15); |
| const int prob = get_prob(p15 << shift, CDF_PROB_TOP); |
| assert(prob >= 128); |
| return av1_prob_cost[prob - 128] + av1_cost_literal(shift); |
| } |
| -------------------------------------------------------------- |
| """ |
| if p15 > (CDF_PROB_TOP - 1): |
| p15 = CDF_PROB_TOP - 1 |
| if p15 < 1: |
| p15 = 1 |
| msb = flog2(int(p15)) # or int(math.floor(math.log2(p15))) using math |
| shift = CDF_PROB_BITS - 1 - msb |
| prob_scaled = p15 * (2**shift) |
| prob = get_prob(prob_scaled, CDF_PROB_TOP) |
| if prob < 128: |
| raise ValueError( |
| f"Normalized probability value is less than 128 (prob={prob},msb={msb},prob_scaled={prob_scaled})" |
| ) |
| return AV1_PROB_COST[prob - 128] + cost_literal_av1(shift) |
| |
| |
| def pmf2cdfinv_av1(pmf): |
| """converts pmf to cdf-inverse""" |
| cdf = CDF_INIT_TOP - np.cumsum(pmf) |
| return cdf |
| |
| |
| def cdfinv2pmf_av1(cdf_inv): |
| """converts cdf-inverse to pmf""" |
| cdf = np.insert(cdf_inv, 0, CDF_INIT_TOP) |
| pmf = np.diff(CDF_INIT_TOP - cdf) |
| return pmf |
| |
| |
| def pmf2cdf_av1(pmf): |
| """converts pmf to cdf""" |
| return CDF_INIT_TOP - pmf2cdfinv_av1(pmf) |
| |
| |
| def cdf2pmf_av1(cdf): |
| """converts cdf to pmf""" |
| cdf_ext = np.insert(cdf, 0, 0) |
| pmf = np.diff(cdf_ext) |
| return pmf |
| |
| |
| def count2cdf_av1(value_count): |
| """Python implementation of the following C code from AVM codec: |
| -------------------------------------------------------------- |
| static void counts_to_cdf(const aom_count_type *counts, aom_cdf_prob *cdf, int modes) { |
| int64_t csum[CDF_MAX_SIZE]; |
| assert(modes <= CDF_MAX_SIZE); |
| |
| csum[0] = counts[0] + 1; |
| for (int i = 1; i < modes; ++i) csum[i] = counts[i] + 1 + csum[i - 1]; |
| |
| for (int i = 0; i < modes; ++i) fprintf(logfile, "%d ", counts[i]); |
| fprintf(logfile, "\n"); |
| |
| int64_t sum = csum[modes - 1]; |
| const int64_t round_shift = sum >> 1; |
| for (int i = 0; i < modes; ++i) { |
| cdf[i] = (csum[i] * CDF_PROB_TOP + round_shift) / sum; |
| cdf[i] = AOMMIN(cdf[i], CDF_PROB_TOP - (modes - 1 + i) * 4); |
| cdf[i] = (i == 0) ? AOMMAX(cdf[i], 4) : AOMMAX(cdf[i], cdf[i - 1] + 4); |
| } |
| } |
| -------------------------------------------------------------- |
| """ |
| value_count = value_count + 1 |
| cdf_count = np.cumsum(value_count) |
| total_count = value_count.sum() |
| round_shift = total_count // 2 |
| nsymb = len(cdf_count) |
| cdf = np.zeros(nsymb).astype(int) |
| for i in range(nsymb): |
| cdf[i] = int((cdf_count[i] * CDF_INIT_TOP + round_shift) / total_count) |
| cdf[i] = min(cdf[i], CDF_INIT_TOP - (nsymb - 1 + i) * 4) |
| if i == 0: |
| cdf[i] = max(cdf[i], 4) |
| else: |
| cdf[i] = max(cdf[i], cdf[i - 1] + 4) |
| return cdf |
