X-Git-Url: https://git.sesse.net/?a=blobdiff_plain;f=qdc.cpp;h=4bb3e9e75c3696eec5b60ed7a976fbefb212d912;hb=3fb87c6b953be3382cd216c74ff6aa025c8eaa2a;hp=a0fe2e7b1a71a36f8f9fe0682cf710218669ca85;hpb=8989d0f3ce3d2c5074aeb5be4297fda30d3efc67;p=narabu diff --git a/qdc.cpp b/qdc.cpp index a0fe2e7..4bb3e9e 100644 --- a/qdc.cpp +++ b/qdc.cpp @@ -5,17 +5,34 @@ #include #include -//#include "ryg_rans/rans64.h" -#include "ryg_rans/rans_byte.h" +#include "ryg_rans/rans64.h" +//#include "ryg_rans/rans_byte.h" #include "ryg_rans/renormalize.h" +#include #include +#include +#include #include +#include #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; @@ -137,21 +154,43 @@ void SymbolStats::normalize_freqs(uint32_t target_total) calc_cost, (calc_cost - ideal_cost) / 8.0, total_loss / 8.0, total_loss_with_dp / 8.0); } -SymbolStats stats[64]; +SymbolStats stats[128]; -int pick_stats_for(int y, int x) +#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 x, int y, bool is_chroma) { - //return 0; - //return std::min(hypot(x, y), 7); - return std::min(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 } @@ -159,11 +198,7 @@ int pick_stats_for(int y, int x) void write_varint(int x, FILE *fp) { - while (x >= 128) { - putc((x & 0x7f) | 0x80, fp); - x >>= 7; - } - putc(x, fp); + fwrite(&x, sizeof(x), 1, fp); } class RansEncoder { @@ -177,22 +212,22 @@ public: 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); - RansEncSymbolInit(&esyms[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); + Rans64EncSymbolInit(&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() { out_end = out_buf.get() + out_max_size; ptr = out_end; // *end* of output buffer - RansEncInit(&rans); + Rans64EncInit(&rans); } uint32_t save_block(FILE *codedfp) // Returns number of bytes. { - RansEncFlush(&rans, &ptr); + Rans64EncFlush(&rans, (uint32_t **)&ptr); //printf("post-flush = %08x\n", rans); uint32_t num_rans_bytes = out_end - ptr; @@ -222,23 +257,23 @@ public: 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; } void encode_coeff(short signed_k) { - //printf("encoding coeff %d (sym %d), rans before encoding = %08x\n", signed_k, ((abs(signed_k) - 1) & 255), rans); + //printf("encoding coeff %d (sym %d), rans before encoding = %016lx\n", signed_k, ((abs(signed_k) - 1) & 255), rans); unsigned short k = abs(signed_k); if (k >= ESCAPE_LIMIT) { // Put the coefficient as a 1/(2^12) symbol _before_ // the 255 coefficient, since the decoder will read the // 255 coefficient first. - RansEncPut(&rans, &ptr, k, 1, prob_bits); + Rans64EncPut(&rans, (uint32_t **)&ptr, k, 1, prob_bits); k = ESCAPE_LIMIT; } - RansEncPutSymbol(&rans, &ptr, &esyms[(k - 1) & 255]); + Rans64EncPutSymbol(&rans, (uint32_t **)&ptr, &esyms[(k - 1) & (NUM_SYMS - 1)], prob_bits + 1); if (signed_k < 0) { rans += sign_bias; } @@ -251,8 +286,8 @@ private: unique_ptr out_buf; uint8_t *out_end; uint8_t *ptr; - RansState rans; - RansEncSymbol esyms[NUM_SYMS]; + Rans64State rans; + Rans64EncSymbol esyms[NUM_SYMS]; uint32_t sign_bias; uint32_t last_block = 0; // Not a valid 4-byte rANS block (?) @@ -348,6 +383,114 @@ void convert_ycbcr() } } +#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 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) @@ -574,10 +717,35 @@ int main(int argc, char **argv) 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; } } @@ -620,9 +788,8 @@ int main(int argc, char **argv) } 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) { @@ -632,7 +799,7 @@ int main(int argc, char **argv) k = ESCAPE_LIMIT; extra_bits += 12; // escape this one } - ++s_luma.freqs[(k - 1) & 255]; + ++s_luma.freqs[(k - 1) & (NUM_SYMS - 1)]; } } // Chroma @@ -648,17 +815,26 @@ int main(int argc, char **argv) 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"); @@ -685,7 +861,7 @@ int main(int argc, char **argv) // 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 @@ -694,21 +870,20 @@ int main(int argc, char **argv) // 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); @@ -730,24 +905,23 @@ int main(int argc, char **argv) // 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); } @@ -756,24 +930,23 @@ int main(int argc, char **argv) // 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); }