#include "ryg_rans/rans_byte.h"
#include "ryg_rans/renormalize.h"
+#include <algorithm>
#include <memory>
+#include <numeric>
+#include <random>
#include <vector>
+#include <unordered_map>
#define WIDTH 1280
#define HEIGHT 720
+#define WIDTH_BLOCKS (WIDTH/8)
+#define WIDTH_BLOCKS_CHROMA (WIDTH/16)
+#define HEIGHT_BLOCKS (HEIGHT/8)
+#define NUM_BLOCKS (WIDTH_BLOCKS * HEIGHT_BLOCKS)
+#define NUM_BLOCKS_CHROMA (WIDTH_BLOCKS_CHROMA * HEIGHT_BLOCKS)
+
#define NUM_SYMS 256
#define ESCAPE_LIMIT (NUM_SYMS - 1)
+#define BLOCKS_PER_STREAM 320
+
+// If you set this to 1, the program will try to optimize the placement
+// of coefficients to rANS probability distributions. This is randomized,
+// so you might want to run it a few times.
+#define FIND_OPTIMAL_STREAM_ASSIGNMENT 0
+#define NUM_CLUSTERS 4
static constexpr uint32_t prob_bits = 12;
static constexpr uint32_t prob_scale = 1 << prob_bits;
calc_cost, (calc_cost - ideal_cost) / 8.0, total_loss / 8.0, total_loss_with_dp / 8.0);
}
-SymbolStats stats[64];
+SymbolStats stats[128];
+
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+// Distance from one stream to the other, based on a hacked-up K-L divergence.
+float kl_dist[64][64];
+#endif
+
+const int luma_mapping[64] = {
+ 0, 0, 1, 1, 2, 2, 3, 3,
+ 0, 0, 1, 2, 2, 2, 3, 3,
+ 1, 1, 2, 2, 2, 3, 3, 3,
+ 1, 1, 2, 2, 2, 3, 3, 3,
+ 1, 2, 2, 2, 2, 3, 3, 3,
+ 2, 2, 2, 2, 3, 3, 3, 3,
+ 2, 2, 3, 3, 3, 3, 3, 3,
+ 3, 3, 3, 3, 3, 3, 3, 3,
+};
+const int chroma_mapping[64] = {
+ 0, 1, 1, 2, 2, 2, 3, 3,
+ 1, 1, 2, 2, 2, 3, 3, 3,
+ 2, 2, 2, 2, 3, 3, 3, 3,
+ 2, 2, 2, 3, 3, 3, 3, 3,
+ 2, 3, 3, 3, 3, 3, 3, 3,
+ 3, 3, 3, 3, 3, 3, 3, 3,
+ 3, 3, 3, 3, 3, 3, 3, 3,
+ 3, 3, 3, 3, 3, 3, 3, 3,
+};
-int pick_stats_for(int y, int x)
+int pick_stats_for(int x, int y, bool is_chroma)
{
- //return 0;
- //return std::min<int>(hypot(x, y), 7);
- return std::min<int>(x + y, 7);
- //if (x + y >= 7) return 7;
- //return x + y;
- //return y * 8 + x;
-#if 0
- if (y == 0 && x == 0) {
- return 0;
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+ return y * 8 + x + is_chroma * 64;
+#else
+ if (is_chroma) {
+ return chroma_mapping[y * 8 + x] + 4;
} else {
- return 1;
+ return luma_mapping[y * 8 + x];
}
#endif
}
void init_prob(SymbolStats &s)
{
for (int i = 0; i < NUM_SYMS; i++) {
- printf("%d: cumfreqs=%d freqs=%d prob_bits=%d\n", i, s.cum_freqs[i], s.freqs[i], prob_bits + 1);
+ //printf("%d: cumfreqs=%d freqs=%d prob_bits=%d\n", i, s.cum_freqs[i], s.freqs[i], prob_bits + 1);
RansEncSymbolInit(&esyms[i], s.cum_freqs[i], s.freqs[i], prob_bits + 1);
}
- sign_bias = s.cum_freqs[256];
+ sign_bias = s.cum_freqs[NUM_SYMS];
}
void clear()
//printf("post-flush = %08x\n", rans);
uint32_t num_rans_bytes = out_end - ptr;
-#if 0
- if (num_rans_bytes == 4) {
- uint32_t block;
- memcpy(&block, ptr, 4);
-
- if (block == last_block) {
- write_varint(0, codedfp);
- clear();
- return 1;
- }
-
- last_block = block;
+ if (num_rans_bytes == last_block.size() &&
+ memcmp(last_block.data(), ptr, last_block.size()) == 0) {
+ write_varint(0, codedfp);
+ clear();
+ return 1;
} else {
- last_block = 0;
+ last_block = string((const char *)ptr, num_rans_bytes);
}
-#endif
write_varint(num_rans_bytes, codedfp);
//fwrite(&num_rans_bytes, 1, 4, codedfp);
clear();
- printf("Saving block: %d rANS bytes\n", num_rans_bytes);
+ //printf("Saving block: %d rANS bytes\n", num_rans_bytes);
return num_rans_bytes;
//return num_rans_bytes;
}
RansEncPut(&rans, &ptr, k, 1, prob_bits);
k = ESCAPE_LIMIT;
}
- RansEncPutSymbol(&rans, &ptr, &esyms[(k - 1) & 255]);
+ RansEncPutSymbol(&rans, &ptr, &esyms[(k - 1) & (NUM_SYMS - 1)]);
if (signed_k < 0) {
rans += sign_bias;
}
RansEncSymbol esyms[NUM_SYMS];
uint32_t sign_bias;
- uint32_t last_block = 0; // Not a valid 4-byte rANS block (?)
+ std::string last_block;
};
static constexpr int dc_scalefac = 8; // Matches the FDCT's gain.
}
}
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+double find_best_assignment(const int *medoids, int *assignment)
+{
+ double current_score = 0.0;
+ for (int i = 0; i < 64; ++i) {
+ int best_medoid = medoids[0];
+ float best_medoid_score = kl_dist[i][medoids[0]];
+ for (int j = 1; j < NUM_CLUSTERS; ++j) {
+ if (kl_dist[i][medoids[j]] < best_medoid_score) {
+ best_medoid = medoids[j];
+ best_medoid_score = kl_dist[i][medoids[j]];
+ }
+ }
+ assignment[i] = best_medoid;
+ current_score += best_medoid_score;
+ }
+ return current_score;
+}
+
+void find_optimal_stream_assignment(int base)
+{
+ double inv_sum[64];
+ for (unsigned i = 0; i < 64; ++i) {
+ double s = 0.0;
+ for (unsigned k = 0; k < NUM_SYMS; ++k) {
+ s += stats[i + base].freqs[k] + 0.5;
+ }
+ inv_sum[i] = 1.0 / s;
+ }
+
+ for (unsigned i = 0; i < 64; ++i) {
+ for (unsigned j = 0; j < 64; ++j) {
+ double d = 0.0;
+ for (unsigned k = 0; k < NUM_SYMS; ++k) {
+ double p1 = (stats[i + base].freqs[k] + 0.5) * inv_sum[i];
+ double p2 = (stats[j + base].freqs[k] + 0.5) * inv_sum[j];
+
+ // K-L divergence is asymmetric; this is a hack.
+ d += p1 * log(p1 / p2);
+ d += p2 * log(p2 / p1);
+ }
+ kl_dist[i][j] = d;
+ //printf("%.3f ", d);
+ }
+ //printf("\n");
+ }
+
+ // k-medoids init
+ int medoids[64]; // only the first NUM_CLUSTERS matter
+ bool is_medoid[64] = { false };
+ std::iota(medoids, medoids + 64, 0);
+ std::random_device rd;
+ std::mt19937 g(rd());
+ std::shuffle(medoids, medoids + 64, g);
+ for (int i = 0; i < NUM_CLUSTERS; ++i) {
+ printf("%d ", medoids[i]);
+ is_medoid[medoids[i]] = true;
+ }
+ printf("\n");
+
+ // assign each data point to the closest medoid
+ int assignment[64];
+ double current_score = find_best_assignment(medoids, assignment);
+
+ for (int i = 0; i < 1000; ++i) {
+ printf("iter %d\n", i);
+ bool any_changed = false;
+ for (int m = 0; m < NUM_CLUSTERS; ++m) {
+ for (int o = 0; o < 64; ++o) {
+ if (is_medoid[o]) continue;
+ int old_medoid = medoids[m];
+ medoids[m] = o;
+
+ int new_assignment[64];
+ double candidate_score = find_best_assignment(medoids, new_assignment);
+
+ if (candidate_score < current_score) {
+ current_score = candidate_score;
+ memcpy(assignment, new_assignment, sizeof(assignment));
+
+ is_medoid[old_medoid] = false;
+ is_medoid[medoids[m]] = true;
+ printf("%f: ", current_score);
+ for (int i = 0; i < 64; ++i) {
+ printf("%d ", assignment[i]);
+ }
+ printf("\n");
+ any_changed = true;
+ } else {
+ medoids[m] = old_medoid;
+ }
+ }
+ }
+ if (!any_changed) break;
+ }
+ printf("\n");
+ std::unordered_map<int, int> rmap;
+ for (int i = 0; i < 64; ++i) {
+ if (i % 8 == 0) printf("\n");
+ if (!rmap.count(assignment[i])) {
+ rmap.emplace(assignment[i], rmap.size());
+ }
+ printf("%d, ", rmap[assignment[i]]);
+ }
+ printf("\n");
+}
+#endif
+
int main(int argc, char **argv)
{
if (argc >= 2)
//double last_cb_cfl_fac = 0.0;
//double last_cr_cfl_fac = 0.0;
+ int max_val_x[8] = {0}, min_val_x[8] = {0};
+ int max_val_y[8] = {0}, min_val_y[8] = {0};
+
// DCT and quantize luma
for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
for (unsigned xb = 0; xb < WIDTH; xb += 8) {
int k = quantize(in_y[coeff_idx], coeff_idx);
coeff_y[(yb + y) * WIDTH + (xb + x)] = k;
+ max_val_x[x] = std::max(max_val_x[x], k);
+ min_val_x[x] = std::min(min_val_x[x], k);
+ max_val_y[y] = std::max(max_val_y[y], k);
+ min_val_y[y] = std::min(min_val_y[y], k);
+
// Store back for reconstruction / PSNR calculation
in_y[coeff_idx] = unquantize(k, coeff_idx);
}
chroma_energy / (WIDTH * HEIGHT), chroma_energy_pred / (WIDTH * HEIGHT));
#endif
- // DC coefficient pred from the right to left
- for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
- for (unsigned xb = 0; xb < WIDTH - 8; xb += 8) {
- coeff_y[yb * WIDTH + xb] -= coeff_y[yb * WIDTH + (xb + 8)];
+ // DC coefficient pred from the right to left (within each slice)
+ for (unsigned block_idx = 0; block_idx < NUM_BLOCKS; block_idx += BLOCKS_PER_STREAM) {
+ int prev_k = 128;
+
+ for (unsigned subblock_idx = BLOCKS_PER_STREAM; subblock_idx --> 0; ) {
+ unsigned yb = (block_idx + subblock_idx) / WIDTH_BLOCKS;
+ unsigned xb = (block_idx + subblock_idx) % WIDTH_BLOCKS;
+ int k = coeff_y[(yb * 8) * WIDTH + (xb * 8)];
+
+ coeff_y[(yb * 8) * WIDTH + (xb * 8)] = k - prev_k;
+
+ prev_k = k;
+ }
+ }
+ for (unsigned block_idx = 0; block_idx < NUM_BLOCKS_CHROMA; block_idx += BLOCKS_PER_STREAM) {
+ int prev_k_cb = 0;
+ int prev_k_cr = 0;
+
+ for (unsigned subblock_idx = BLOCKS_PER_STREAM; subblock_idx --> 0; ) {
+ unsigned yb = (block_idx + subblock_idx) / WIDTH_BLOCKS_CHROMA;
+ unsigned xb = (block_idx + subblock_idx) % WIDTH_BLOCKS_CHROMA;
+ int k_cb = coeff_cb[(yb * 8) * WIDTH/2 + (xb * 8)];
+ int k_cr = coeff_cr[(yb * 8) * WIDTH/2 + (xb * 8)];
+
+ coeff_cb[(yb * 8) * WIDTH/2 + (xb * 8)] = k_cb - prev_k_cb;
+ coeff_cr[(yb * 8) * WIDTH/2 + (xb * 8)] = k_cr - prev_k_cr;
+
+ prev_k_cb = k_cb;
+ prev_k_cr = k_cr;
}
}
}
for (unsigned y = 0; y < 8; ++y) {
for (unsigned x = 0; x < 8; ++x) {
- SymbolStats &s_luma = stats[pick_stats_for(x, y)];
- SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8]; // HACK
- //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+ SymbolStats &s_luma = stats[pick_stats_for(x, y, false)];
+ SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
// Luma
for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
k = ESCAPE_LIMIT;
extra_bits += 12; // escape this one
}
- ++s_luma.freqs[(k - 1) & 255];
+ ++s_luma.freqs[(k - 1) & (NUM_SYMS - 1)];
}
}
// Chroma
k_cr = ESCAPE_LIMIT;
extra_bits += 12; // escape this one
}
- ++s_chroma.freqs[(k_cb - 1) & 255];
- ++s_chroma.freqs[(k_cr - 1) & 255];
+ ++s_chroma.freqs[(k_cb - 1) & (NUM_SYMS - 1)];
+ ++s_chroma.freqs[(k_cr - 1) & (NUM_SYMS - 1)];
}
}
}
}
+
+#if FIND_OPTIMAL_STREAM_ASSIGNMENT
+ printf("Luma:\n");
+ find_optimal_stream_assignment(0);
+ printf("Chroma:\n");
+ find_optimal_stream_assignment(64);
+ exit(0);
+#endif
+
for (unsigned i = 0; i < 64; ++i) {
- stats[i].freqs[255] /= 2; // zero, has no sign bits (yes, this is trickery)
+ stats[i].freqs[NUM_SYMS - 1] /= 2; // zero, has no sign bits (yes, this is trickery)
stats[i].normalize_freqs(prob_scale);
- stats[i].cum_freqs[256] += stats[i].freqs[255];
- stats[i].freqs[255] *= 2;
+ stats[i].cum_freqs[NUM_SYMS] += stats[i].freqs[NUM_SYMS - 1];
+ stats[i].freqs[NUM_SYMS - 1] *= 2;
}
FILE *codedfp = fopen("coded.dat", "wb");
// Luma
for (unsigned y = 0; y < 8; ++y) {
for (unsigned x = 0; x < 8; ++x) {
- SymbolStats &s_luma = stats[pick_stats_for(x, y)];
+ SymbolStats &s_luma = stats[pick_stats_for(x, y, false)];
rans_encoder.init_prob(s_luma);
// Luma
// need to reverse later
rans_encoder.clear();
size_t num_bytes = 0;
- for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
- for (unsigned xb = 0; xb < WIDTH; xb += 8) {
- int k = coeff_y[(yb + y) * WIDTH + (xb + x)];
- //printf("encoding coeff %d xb,yb=%d,%d: %d\n", y*8+x, xb, yb, k);
- rans_encoder.encode_coeff(k);
- }
- if (yb % 16 == 8) {
+ for (unsigned block_idx = 0; block_idx < NUM_BLOCKS; ++block_idx) {
+ unsigned yb = block_idx / WIDTH_BLOCKS;
+ unsigned xb = block_idx % WIDTH_BLOCKS;
+
+ int k = coeff_y[(yb * 8 + y) * WIDTH + (xb * 8 + x)];
+ //printf("encoding coeff %d xb,yb=%d,%d: %d\n", y*8+x, xb, yb, k);
+ rans_encoder.encode_coeff(k);
+
+ if (block_idx % BLOCKS_PER_STREAM == (BLOCKS_PER_STREAM - 1) || block_idx == NUM_BLOCKS - 1) {
int l = rans_encoder.save_block(codedfp);
num_bytes += l;
lens.push_back(l);
}
}
- if (HEIGHT % 16 != 0) {
- num_bytes += rans_encoder.save_block(codedfp);
- }
tot_bytes += num_bytes;
printf("coeff %d Y': %ld bytes\n", y * 8 + x, num_bytes);
// Cb
for (unsigned y = 0; y < 8; ++y) {
for (unsigned x = 0; x < 8; ++x) {
- SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8];
- //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+ SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
rans_encoder.init_prob(s_chroma);
rans_encoder.clear();
size_t num_bytes = 0;
- for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
- for (unsigned xb = 0; xb < WIDTH/2; xb += 8) {
- int k = coeff_cb[(yb + y) * WIDTH/2 + (xb + x)];
- rans_encoder.encode_coeff(k);
- }
- if (yb % 16 == 8) {
+ for (unsigned block_idx = 0; block_idx < NUM_BLOCKS_CHROMA; ++block_idx) {
+ unsigned yb = block_idx / WIDTH_BLOCKS_CHROMA;
+ unsigned xb = block_idx % WIDTH_BLOCKS_CHROMA;
+
+ int k = coeff_cb[(yb * 8 + y) * WIDTH/2 + (xb * 8 + x)];
+ //printf("encoding coeff %d xb,yb=%d,%d: %d\n", y*8+x, xb, yb, k);
+ rans_encoder.encode_coeff(k);
+
+ if (block_idx % BLOCKS_PER_STREAM == (BLOCKS_PER_STREAM - 1) || block_idx == NUM_BLOCKS - 1) {
num_bytes += rans_encoder.save_block(codedfp);
}
}
- if (HEIGHT % 16 != 0) {
- num_bytes += rans_encoder.save_block(codedfp);
- }
tot_bytes += num_bytes;
printf("coeff %d Cb: %ld bytes\n", y * 8 + x, num_bytes);
}
// Cr
for (unsigned y = 0; y < 8; ++y) {
for (unsigned x = 0; x < 8; ++x) {
- SymbolStats &s_chroma = stats[pick_stats_for(x, y) + 8];
- //SymbolStats &s_chroma = stats[pick_stats_for(x, y)];
+ SymbolStats &s_chroma = stats[pick_stats_for(x, y, true)];
rans_encoder.init_prob(s_chroma);
rans_encoder.clear();
size_t num_bytes = 0;
- for (unsigned yb = 0; yb < HEIGHT; yb += 8) {
- for (unsigned xb = 0; xb < WIDTH/2; xb += 8) {
- int k = coeff_cr[(yb + y) * WIDTH/2 + (xb + x)];
- rans_encoder.encode_coeff(k);
- }
- if (yb % 16 == 8) {
+ for (unsigned block_idx = 0; block_idx < NUM_BLOCKS_CHROMA; ++block_idx) {
+ unsigned yb = block_idx / WIDTH_BLOCKS_CHROMA;
+ unsigned xb = block_idx % WIDTH_BLOCKS_CHROMA;
+
+ int k = coeff_cr[(yb * 8 + y) * WIDTH/2 + (xb * 8 + x)];
+ //printf("encoding coeff %d xb,yb=%d,%d: %d\n", y*8+x, xb, yb, k);
+ rans_encoder.encode_coeff(k);
+
+ if (block_idx % BLOCKS_PER_STREAM == (BLOCKS_PER_STREAM - 1) || block_idx == NUM_BLOCKS - 1) {
num_bytes += rans_encoder.save_block(codedfp);
}
}
- if (HEIGHT % 16 != 0) {
- num_bytes += rans_encoder.save_block(codedfp);
- }
tot_bytes += num_bytes;
printf("coeff %d Cr: %ld bytes\n", y * 8 + x, num_bytes);
}
extra_bits,
extra_bits / 8,
tot_bytes);
+
+#if 0
+ printf("Max coefficient ranges (as a function of x):\n\n");
+ for (unsigned x = 0; x < 8; ++x) {
+ int range = std::max(max_val_x[x], -min_val_x[x]);
+ printf(" [%4d, %4d] (%.2f bits)\n", min_val_x[x], max_val_x[x], log2(range * 2 + 1));
+ }
+
+ printf("Max coefficient ranges (as a function of y):\n\n");
+ for (unsigned y = 0; y < 8; ++y) {
+ int range = std::max(max_val_y[y], -min_val_y[y]);
+ printf(" [%4d, %4d] (%.2f bits)\n", min_val_y[y], max_val_y[y], log2(range * 2 + 1));
+ }
+#endif
}