/src/aom/av1/encoder/palette.c
Line | Count | Source (jump to first uncovered line) |
1 | | /* |
2 | | * Copyright (c) 2016, Alliance for Open Media. All rights reserved |
3 | | * |
4 | | * This source code is subject to the terms of the BSD 2 Clause License and |
5 | | * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License |
6 | | * was not distributed with this source code in the LICENSE file, you can |
7 | | * obtain it at www.aomedia.org/license/software. If the Alliance for Open |
8 | | * Media Patent License 1.0 was not distributed with this source code in the |
9 | | * PATENTS file, you can obtain it at www.aomedia.org/license/patent. |
10 | | */ |
11 | | |
12 | | #include <math.h> |
13 | | #include <stdlib.h> |
14 | | |
15 | | #include "av1/common/pred_common.h" |
16 | | |
17 | | #include "av1/encoder/block.h" |
18 | | #include "av1/encoder/cost.h" |
19 | | #include "av1/encoder/encoder.h" |
20 | | #include "av1/encoder/intra_mode_search.h" |
21 | | #include "av1/encoder/intra_mode_search_utils.h" |
22 | | #include "av1/encoder/palette.h" |
23 | | #include "av1/encoder/random.h" |
24 | | #include "av1/encoder/rdopt_utils.h" |
25 | | #include "av1/encoder/tx_search.h" |
26 | | |
27 | 0 | #define AV1_K_MEANS_DIM 1 |
28 | | #include "av1/encoder/k_means_template.h" |
29 | | #undef AV1_K_MEANS_DIM |
30 | 0 | #define AV1_K_MEANS_DIM 2 |
31 | | #include "av1/encoder/k_means_template.h" |
32 | | #undef AV1_K_MEANS_DIM |
33 | | |
34 | 0 | static int int_comparer(const void *a, const void *b) { |
35 | 0 | return (*(int *)a - *(int *)b); |
36 | 0 | } |
37 | | |
38 | 0 | int av1_remove_duplicates(int *centroids, int num_centroids) { |
39 | 0 | int num_unique; // number of unique centroids |
40 | 0 | int i; |
41 | 0 | qsort(centroids, num_centroids, sizeof(*centroids), int_comparer); |
42 | | // Remove duplicates. |
43 | 0 | num_unique = 1; |
44 | 0 | for (i = 1; i < num_centroids; ++i) { |
45 | 0 | if (centroids[i] != centroids[i - 1]) { // found a new unique centroid |
46 | 0 | centroids[num_unique++] = centroids[i]; |
47 | 0 | } |
48 | 0 | } |
49 | 0 | return num_unique; |
50 | 0 | } |
51 | | |
52 | | static int delta_encode_cost(const int *colors, int num, int bit_depth, |
53 | 0 | int min_val) { |
54 | 0 | if (num <= 0) return 0; |
55 | 0 | int bits_cost = bit_depth; |
56 | 0 | if (num == 1) return bits_cost; |
57 | 0 | bits_cost += 2; |
58 | 0 | int max_delta = 0; |
59 | 0 | int deltas[PALETTE_MAX_SIZE]; |
60 | 0 | const int min_bits = bit_depth - 3; |
61 | 0 | for (int i = 1; i < num; ++i) { |
62 | 0 | const int delta = colors[i] - colors[i - 1]; |
63 | 0 | deltas[i - 1] = delta; |
64 | 0 | assert(delta >= min_val); |
65 | 0 | if (delta > max_delta) max_delta = delta; |
66 | 0 | } |
67 | 0 | int bits_per_delta = AOMMAX(av1_ceil_log2(max_delta + 1 - min_val), min_bits); |
68 | 0 | assert(bits_per_delta <= bit_depth); |
69 | 0 | int range = (1 << bit_depth) - colors[0] - min_val; |
70 | 0 | for (int i = 0; i < num - 1; ++i) { |
71 | 0 | bits_cost += bits_per_delta; |
72 | 0 | range -= deltas[i]; |
73 | 0 | bits_per_delta = AOMMIN(bits_per_delta, av1_ceil_log2(range)); |
74 | 0 | } |
75 | 0 | return bits_cost; |
76 | 0 | } |
77 | | |
78 | | int av1_index_color_cache(const uint16_t *color_cache, int n_cache, |
79 | | const uint16_t *colors, int n_colors, |
80 | 0 | uint8_t *cache_color_found, int *out_cache_colors) { |
81 | 0 | if (n_cache <= 0) { |
82 | 0 | for (int i = 0; i < n_colors; ++i) out_cache_colors[i] = colors[i]; |
83 | 0 | return n_colors; |
84 | 0 | } |
85 | 0 | memset(cache_color_found, 0, n_cache * sizeof(*cache_color_found)); |
86 | 0 | int n_in_cache = 0; |
87 | 0 | int in_cache_flags[PALETTE_MAX_SIZE]; |
88 | 0 | memset(in_cache_flags, 0, sizeof(in_cache_flags)); |
89 | 0 | for (int i = 0; i < n_cache && n_in_cache < n_colors; ++i) { |
90 | 0 | for (int j = 0; j < n_colors; ++j) { |
91 | 0 | if (colors[j] == color_cache[i]) { |
92 | 0 | in_cache_flags[j] = 1; |
93 | 0 | cache_color_found[i] = 1; |
94 | 0 | ++n_in_cache; |
95 | 0 | break; |
96 | 0 | } |
97 | 0 | } |
98 | 0 | } |
99 | 0 | int j = 0; |
100 | 0 | for (int i = 0; i < n_colors; ++i) |
101 | 0 | if (!in_cache_flags[i]) out_cache_colors[j++] = colors[i]; |
102 | 0 | assert(j == n_colors - n_in_cache); |
103 | 0 | return j; |
104 | 0 | } |
105 | | |
106 | | int av1_get_palette_delta_bits_v(const PALETTE_MODE_INFO *const pmi, |
107 | | int bit_depth, int *zero_count, |
108 | 0 | int *min_bits) { |
109 | 0 | const int n = pmi->palette_size[1]; |
110 | 0 | const int max_val = 1 << bit_depth; |
111 | 0 | int max_d = 0; |
112 | 0 | *min_bits = bit_depth - 4; |
113 | 0 | *zero_count = 0; |
114 | 0 | for (int i = 1; i < n; ++i) { |
115 | 0 | const int delta = pmi->palette_colors[2 * PALETTE_MAX_SIZE + i] - |
116 | 0 | pmi->palette_colors[2 * PALETTE_MAX_SIZE + i - 1]; |
117 | 0 | const int v = abs(delta); |
118 | 0 | const int d = AOMMIN(v, max_val - v); |
119 | 0 | if (d > max_d) max_d = d; |
120 | 0 | if (d == 0) ++(*zero_count); |
121 | 0 | } |
122 | 0 | return AOMMAX(av1_ceil_log2(max_d + 1), *min_bits); |
123 | 0 | } |
124 | | |
125 | | int av1_palette_color_cost_y(const PALETTE_MODE_INFO *const pmi, |
126 | | const uint16_t *color_cache, int n_cache, |
127 | 0 | int bit_depth) { |
128 | 0 | const int n = pmi->palette_size[0]; |
129 | 0 | int out_cache_colors[PALETTE_MAX_SIZE]; |
130 | 0 | uint8_t cache_color_found[2 * PALETTE_MAX_SIZE]; |
131 | 0 | const int n_out_cache = |
132 | 0 | av1_index_color_cache(color_cache, n_cache, pmi->palette_colors, n, |
133 | 0 | cache_color_found, out_cache_colors); |
134 | 0 | const int total_bits = |
135 | 0 | n_cache + delta_encode_cost(out_cache_colors, n_out_cache, bit_depth, 1); |
136 | 0 | return av1_cost_literal(total_bits); |
137 | 0 | } |
138 | | |
139 | | int av1_palette_color_cost_uv(const PALETTE_MODE_INFO *const pmi, |
140 | | const uint16_t *color_cache, int n_cache, |
141 | 0 | int bit_depth) { |
142 | 0 | const int n = pmi->palette_size[1]; |
143 | 0 | int total_bits = 0; |
144 | | // U channel palette color cost. |
145 | 0 | int out_cache_colors[PALETTE_MAX_SIZE]; |
146 | 0 | uint8_t cache_color_found[2 * PALETTE_MAX_SIZE]; |
147 | 0 | const int n_out_cache = av1_index_color_cache( |
148 | 0 | color_cache, n_cache, pmi->palette_colors + PALETTE_MAX_SIZE, n, |
149 | 0 | cache_color_found, out_cache_colors); |
150 | 0 | total_bits += |
151 | 0 | n_cache + delta_encode_cost(out_cache_colors, n_out_cache, bit_depth, 0); |
152 | | |
153 | | // V channel palette color cost. |
154 | 0 | int zero_count = 0, min_bits_v = 0; |
155 | 0 | const int bits_v = |
156 | 0 | av1_get_palette_delta_bits_v(pmi, bit_depth, &zero_count, &min_bits_v); |
157 | 0 | const int bits_using_delta = |
158 | 0 | 2 + bit_depth + (bits_v + 1) * (n - 1) - zero_count; |
159 | 0 | const int bits_using_raw = bit_depth * n; |
160 | 0 | total_bits += 1 + AOMMIN(bits_using_delta, bits_using_raw); |
161 | 0 | return av1_cost_literal(total_bits); |
162 | 0 | } |
163 | | |
164 | | // Extends 'color_map' array from 'orig_width x orig_height' to 'new_width x |
165 | | // new_height'. Extra rows and columns are filled in by copying last valid |
166 | | // row/column. |
167 | | static AOM_INLINE void extend_palette_color_map(uint8_t *const color_map, |
168 | | int orig_width, int orig_height, |
169 | 0 | int new_width, int new_height) { |
170 | 0 | int j; |
171 | 0 | assert(new_width >= orig_width); |
172 | 0 | assert(new_height >= orig_height); |
173 | 0 | if (new_width == orig_width && new_height == orig_height) return; |
174 | | |
175 | 0 | for (j = orig_height - 1; j >= 0; --j) { |
176 | 0 | memmove(color_map + j * new_width, color_map + j * orig_width, orig_width); |
177 | | // Copy last column to extra columns. |
178 | 0 | memset(color_map + j * new_width + orig_width, |
179 | 0 | color_map[j * new_width + orig_width - 1], new_width - orig_width); |
180 | 0 | } |
181 | | // Copy last row to extra rows. |
182 | 0 | for (j = orig_height; j < new_height; ++j) { |
183 | 0 | memcpy(color_map + j * new_width, color_map + (orig_height - 1) * new_width, |
184 | 0 | new_width); |
185 | 0 | } |
186 | 0 | } |
187 | | |
188 | | // Bias toward using colors in the cache. |
189 | | // TODO(huisu): Try other schemes to improve compression. |
190 | | static AOM_INLINE void optimize_palette_colors(uint16_t *color_cache, |
191 | | int n_cache, int n_colors, |
192 | | int stride, int *centroids, |
193 | 0 | int bit_depth) { |
194 | 0 | if (n_cache <= 0) return; |
195 | 0 | for (int i = 0; i < n_colors * stride; i += stride) { |
196 | 0 | int min_diff = abs(centroids[i] - (int)color_cache[0]); |
197 | 0 | int idx = 0; |
198 | 0 | for (int j = 1; j < n_cache; ++j) { |
199 | 0 | const int this_diff = abs(centroids[i] - color_cache[j]); |
200 | 0 | if (this_diff < min_diff) { |
201 | 0 | min_diff = this_diff; |
202 | 0 | idx = j; |
203 | 0 | } |
204 | 0 | } |
205 | 0 | const int min_threshold = 4 << (bit_depth - 8); |
206 | 0 | if (min_diff <= min_threshold) centroids[i] = color_cache[idx]; |
207 | 0 | } |
208 | 0 | } |
209 | | |
210 | | /*!\brief Calculate the luma palette cost from a given color palette |
211 | | * |
212 | | * \ingroup palette_mode_search |
213 | | * \callergraph |
214 | | * Given the base colors as specified in centroids[], calculate the RD cost |
215 | | * of palette mode. |
216 | | */ |
217 | | static AOM_INLINE void palette_rd_y( |
218 | | const AV1_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi, |
219 | | BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int *centroids, int n, |
220 | | uint16_t *color_cache, int n_cache, bool do_header_rd_based_gating, |
221 | | MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map, int64_t *best_rd, |
222 | | int *rate, int *rate_tokenonly, int64_t *distortion, int *skippable, |
223 | | int *beat_best_rd, PICK_MODE_CONTEXT *ctx, uint8_t *blk_skip, |
224 | | uint8_t *tx_type_map, int *beat_best_palette_rd, |
225 | 0 | bool *do_header_rd_based_breakout) { |
226 | 0 | if (do_header_rd_based_breakout != NULL) *do_header_rd_based_breakout = false; |
227 | 0 | optimize_palette_colors(color_cache, n_cache, n, 1, centroids, |
228 | 0 | cpi->common.seq_params->bit_depth); |
229 | 0 | const int num_unique_colors = av1_remove_duplicates(centroids, n); |
230 | 0 | if (num_unique_colors < PALETTE_MIN_SIZE) { |
231 | | // Too few unique colors to create a palette. And DC_PRED will work |
232 | | // well for that case anyway. So skip. |
233 | 0 | return; |
234 | 0 | } |
235 | 0 | PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info; |
236 | 0 | if (cpi->common.seq_params->use_highbitdepth) { |
237 | 0 | for (int i = 0; i < num_unique_colors; ++i) { |
238 | 0 | pmi->palette_colors[i] = clip_pixel_highbd( |
239 | 0 | (int)centroids[i], cpi->common.seq_params->bit_depth); |
240 | 0 | } |
241 | 0 | } else { |
242 | 0 | for (int i = 0; i < num_unique_colors; ++i) { |
243 | 0 | pmi->palette_colors[i] = clip_pixel(centroids[i]); |
244 | 0 | } |
245 | 0 | } |
246 | 0 | pmi->palette_size[0] = num_unique_colors; |
247 | 0 | MACROBLOCKD *const xd = &x->e_mbd; |
248 | 0 | uint8_t *const color_map = xd->plane[0].color_index_map; |
249 | 0 | int block_width, block_height, rows, cols; |
250 | 0 | av1_get_block_dimensions(bsize, 0, xd, &block_width, &block_height, &rows, |
251 | 0 | &cols); |
252 | 0 | av1_calc_indices(data, centroids, color_map, rows * cols, num_unique_colors, |
253 | 0 | 1); |
254 | 0 | extend_palette_color_map(color_map, cols, rows, block_width, block_height); |
255 | |
|
256 | 0 | RD_STATS tokenonly_rd_stats; |
257 | 0 | int this_rate; |
258 | |
|
259 | 0 | if (do_header_rd_based_gating) { |
260 | 0 | assert(do_header_rd_based_breakout != NULL); |
261 | 0 | const int palette_mode_rate = |
262 | 0 | intra_mode_info_cost_y(cpi, x, mbmi, bsize, dc_mode_cost); |
263 | 0 | const int64_t header_rd = RDCOST(x->rdmult, palette_mode_rate, 0); |
264 | | // Less aggressive pruning when prune_luma_palette_size_search_level == 1. |
265 | 0 | const int header_rd_shift = |
266 | 0 | (cpi->sf.intra_sf.prune_luma_palette_size_search_level == 1) ? 1 : 0; |
267 | | // Terminate further palette_size search, if the header cost corresponding |
268 | | // to lower palette_size is more than *best_rd << header_rd_shift. This |
269 | | // logic is implemented with a right shift in the LHS to prevent a possible |
270 | | // overflow with the left shift in RHS. |
271 | 0 | if ((header_rd >> header_rd_shift) > *best_rd) { |
272 | 0 | *do_header_rd_based_breakout = true; |
273 | 0 | return; |
274 | 0 | } |
275 | 0 | av1_pick_uniform_tx_size_type_yrd(cpi, x, &tokenonly_rd_stats, bsize, |
276 | 0 | *best_rd); |
277 | 0 | if (tokenonly_rd_stats.rate == INT_MAX) return; |
278 | 0 | this_rate = tokenonly_rd_stats.rate + palette_mode_rate; |
279 | 0 | } else { |
280 | 0 | av1_pick_uniform_tx_size_type_yrd(cpi, x, &tokenonly_rd_stats, bsize, |
281 | 0 | *best_rd); |
282 | 0 | if (tokenonly_rd_stats.rate == INT_MAX) return; |
283 | 0 | this_rate = tokenonly_rd_stats.rate + |
284 | 0 | intra_mode_info_cost_y(cpi, x, mbmi, bsize, dc_mode_cost); |
285 | 0 | } |
286 | | |
287 | 0 | int64_t this_rd = RDCOST(x->rdmult, this_rate, tokenonly_rd_stats.dist); |
288 | 0 | if (!xd->lossless[mbmi->segment_id] && block_signals_txsize(mbmi->bsize)) { |
289 | 0 | tokenonly_rd_stats.rate -= tx_size_cost(x, bsize, mbmi->tx_size); |
290 | 0 | } |
291 | | // Collect mode stats for multiwinner mode processing |
292 | 0 | const int txfm_search_done = 1; |
293 | 0 | store_winner_mode_stats( |
294 | 0 | &cpi->common, x, mbmi, NULL, NULL, NULL, THR_DC, color_map, bsize, |
295 | 0 | this_rd, cpi->sf.winner_mode_sf.multi_winner_mode_type, txfm_search_done); |
296 | 0 | if (this_rd < *best_rd) { |
297 | 0 | *best_rd = this_rd; |
298 | | // Setting beat_best_rd flag because current mode rd is better than best_rd. |
299 | | // This flag need to be updated only for palette evaluation in key frames |
300 | 0 | if (beat_best_rd) *beat_best_rd = 1; |
301 | 0 | memcpy(best_palette_color_map, color_map, |
302 | 0 | block_width * block_height * sizeof(color_map[0])); |
303 | 0 | *best_mbmi = *mbmi; |
304 | 0 | memcpy(blk_skip, x->txfm_search_info.blk_skip, |
305 | 0 | sizeof(x->txfm_search_info.blk_skip[0]) * ctx->num_4x4_blk); |
306 | 0 | av1_copy_array(tx_type_map, xd->tx_type_map, ctx->num_4x4_blk); |
307 | 0 | if (rate) *rate = this_rate; |
308 | 0 | if (rate_tokenonly) *rate_tokenonly = tokenonly_rd_stats.rate; |
309 | 0 | if (distortion) *distortion = tokenonly_rd_stats.dist; |
310 | 0 | if (skippable) *skippable = tokenonly_rd_stats.skip_txfm; |
311 | 0 | if (beat_best_palette_rd) *beat_best_palette_rd = 1; |
312 | 0 | } |
313 | 0 | } |
314 | | |
315 | 0 | static AOM_INLINE int is_iter_over(int curr_idx, int end_idx, int step_size) { |
316 | 0 | assert(step_size != 0); |
317 | 0 | return (step_size > 0) ? curr_idx >= end_idx : curr_idx <= end_idx; |
318 | 0 | } |
319 | | |
320 | | // Performs count-based palette search with number of colors in interval |
321 | | // [start_n, end_n) with step size step_size. If step_size < 0, then end_n can |
322 | | // be less than start_n. Saves the last numbers searched in last_n_searched and |
323 | | // returns the best number of colors found. |
324 | | static AOM_INLINE int perform_top_color_palette_search( |
325 | | const AV1_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi, |
326 | | BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int *top_colors, |
327 | | int start_n, int end_n, int step_size, bool do_header_rd_based_gating, |
328 | | int *last_n_searched, uint16_t *color_cache, int n_cache, |
329 | | MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map, int64_t *best_rd, |
330 | | int *rate, int *rate_tokenonly, int64_t *distortion, int *skippable, |
331 | | int *beat_best_rd, PICK_MODE_CONTEXT *ctx, uint8_t *best_blk_skip, |
332 | 0 | uint8_t *tx_type_map) { |
333 | 0 | int centroids[PALETTE_MAX_SIZE]; |
334 | 0 | int n = start_n; |
335 | 0 | int top_color_winner = end_n; |
336 | | /* clang-format off */ |
337 | 0 | assert(IMPLIES(step_size < 0, start_n > end_n)); |
338 | | /* clang-format on */ |
339 | 0 | assert(IMPLIES(step_size > 0, start_n < end_n)); |
340 | 0 | while (!is_iter_over(n, end_n, step_size)) { |
341 | 0 | int beat_best_palette_rd = 0; |
342 | 0 | bool do_header_rd_based_breakout = false; |
343 | 0 | memcpy(centroids, top_colors, n * sizeof(top_colors[0])); |
344 | 0 | palette_rd_y(cpi, x, mbmi, bsize, dc_mode_cost, data, centroids, n, |
345 | 0 | color_cache, n_cache, do_header_rd_based_gating, best_mbmi, |
346 | 0 | best_palette_color_map, best_rd, rate, rate_tokenonly, |
347 | 0 | distortion, skippable, beat_best_rd, ctx, best_blk_skip, |
348 | 0 | tx_type_map, &beat_best_palette_rd, |
349 | 0 | &do_header_rd_based_breakout); |
350 | 0 | *last_n_searched = n; |
351 | 0 | if (do_header_rd_based_breakout) { |
352 | | // Terminate palette_size search by setting last_n_searched to end_n. |
353 | 0 | *last_n_searched = end_n; |
354 | 0 | break; |
355 | 0 | } |
356 | 0 | if (beat_best_palette_rd) { |
357 | 0 | top_color_winner = n; |
358 | 0 | } else if (cpi->sf.intra_sf.prune_palette_search_level == 2) { |
359 | | // At search level 2, we return immediately if we don't see an improvement |
360 | 0 | return top_color_winner; |
361 | 0 | } |
362 | 0 | n += step_size; |
363 | 0 | } |
364 | 0 | return top_color_winner; |
365 | 0 | } |
366 | | |
367 | | // Performs k-means based palette search with number of colors in interval |
368 | | // [start_n, end_n) with step size step_size. If step_size < 0, then end_n can |
369 | | // be less than start_n. Saves the last numbers searched in last_n_searched and |
370 | | // returns the best number of colors found. |
371 | | static AOM_INLINE int perform_k_means_palette_search( |
372 | | const AV1_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi, |
373 | | BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int lower_bound, |
374 | | int upper_bound, int start_n, int end_n, int step_size, |
375 | | bool do_header_rd_based_gating, int *last_n_searched, uint16_t *color_cache, |
376 | | int n_cache, MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map, |
377 | | int64_t *best_rd, int *rate, int *rate_tokenonly, int64_t *distortion, |
378 | | int *skippable, int *beat_best_rd, PICK_MODE_CONTEXT *ctx, |
379 | | uint8_t *best_blk_skip, uint8_t *tx_type_map, uint8_t *color_map, |
380 | 0 | int data_points) { |
381 | 0 | int centroids[PALETTE_MAX_SIZE]; |
382 | 0 | const int max_itr = 50; |
383 | 0 | int n = start_n; |
384 | 0 | int top_color_winner = end_n; |
385 | | /* clang-format off */ |
386 | 0 | assert(IMPLIES(step_size < 0, start_n > end_n)); |
387 | | /* clang-format on */ |
388 | 0 | assert(IMPLIES(step_size > 0, start_n < end_n)); |
389 | 0 | while (!is_iter_over(n, end_n, step_size)) { |
390 | 0 | int beat_best_palette_rd = 0; |
391 | 0 | bool do_header_rd_based_breakout = false; |
392 | 0 | for (int i = 0; i < n; ++i) { |
393 | 0 | centroids[i] = |
394 | 0 | lower_bound + (2 * i + 1) * (upper_bound - lower_bound) / n / 2; |
395 | 0 | } |
396 | 0 | av1_k_means(data, centroids, color_map, data_points, n, 1, max_itr); |
397 | 0 | palette_rd_y(cpi, x, mbmi, bsize, dc_mode_cost, data, centroids, n, |
398 | 0 | color_cache, n_cache, do_header_rd_based_gating, best_mbmi, |
399 | 0 | best_palette_color_map, best_rd, rate, rate_tokenonly, |
400 | 0 | distortion, skippable, beat_best_rd, ctx, best_blk_skip, |
401 | 0 | tx_type_map, &beat_best_palette_rd, |
402 | 0 | &do_header_rd_based_breakout); |
403 | 0 | *last_n_searched = n; |
404 | 0 | if (do_header_rd_based_breakout) { |
405 | | // Terminate palette_size search by setting last_n_searched to end_n. |
406 | 0 | *last_n_searched = end_n; |
407 | 0 | break; |
408 | 0 | } |
409 | 0 | if (beat_best_palette_rd) { |
410 | 0 | top_color_winner = n; |
411 | 0 | } else if (cpi->sf.intra_sf.prune_palette_search_level == 2) { |
412 | | // At search level 2, we return immediately if we don't see an improvement |
413 | 0 | return top_color_winner; |
414 | 0 | } |
415 | 0 | n += step_size; |
416 | 0 | } |
417 | 0 | return top_color_winner; |
418 | 0 | } |
419 | | |
420 | | // Sets the parameters to search the current number of colors +- 1 |
421 | | static AOM_INLINE void set_stage2_params(int *min_n, int *max_n, int *step_size, |
422 | 0 | int winner, int end_n) { |
423 | | // Set min to winner - 1 unless we are already at the border, then we set it |
424 | | // to winner + 1 |
425 | 0 | *min_n = (winner == PALETTE_MIN_SIZE) ? (PALETTE_MIN_SIZE + 1) |
426 | 0 | : AOMMAX(winner - 1, PALETTE_MIN_SIZE); |
427 | | // Set max to winner + 1 unless we are already at the border, then we set it |
428 | | // to winner - 1 |
429 | 0 | *max_n = |
430 | 0 | (winner == end_n) ? (winner - 1) : AOMMIN(winner + 1, PALETTE_MAX_SIZE); |
431 | | |
432 | | // Set the step size to max_n - min_n so we only search those two values. |
433 | | // If max_n == min_n, then set step_size to 1 to avoid infinite loop later. |
434 | 0 | *step_size = AOMMAX(1, *max_n - *min_n); |
435 | 0 | } |
436 | | |
437 | | static AOM_INLINE void fill_data_and_get_bounds( |
438 | | const uint8_t *src, const int src_stride, const int rows, const int cols, |
439 | 0 | const int is_high_bitdepth, int *data, int *lower_bound, int *upper_bound) { |
440 | 0 | if (is_high_bitdepth) { |
441 | 0 | const uint16_t *src_ptr = CONVERT_TO_SHORTPTR(src); |
442 | 0 | *lower_bound = *upper_bound = src_ptr[0]; |
443 | 0 | for (int r = 0; r < rows; ++r) { |
444 | 0 | for (int c = 0; c < cols; ++c) { |
445 | 0 | const int val = src_ptr[c]; |
446 | 0 | data[c] = val; |
447 | 0 | *lower_bound = AOMMIN(*lower_bound, val); |
448 | 0 | *upper_bound = AOMMAX(*upper_bound, val); |
449 | 0 | } |
450 | 0 | src_ptr += src_stride; |
451 | 0 | data += cols; |
452 | 0 | } |
453 | 0 | return; |
454 | 0 | } |
455 | | |
456 | | // low bit depth |
457 | 0 | *lower_bound = *upper_bound = src[0]; |
458 | 0 | for (int r = 0; r < rows; ++r) { |
459 | 0 | for (int c = 0; c < cols; ++c) { |
460 | 0 | const int val = src[c]; |
461 | 0 | data[c] = val; |
462 | 0 | *lower_bound = AOMMIN(*lower_bound, val); |
463 | 0 | *upper_bound = AOMMAX(*upper_bound, val); |
464 | 0 | } |
465 | 0 | src += src_stride; |
466 | 0 | data += cols; |
467 | 0 | } |
468 | 0 | } |
469 | | |
470 | | void av1_rd_pick_palette_intra_sby( |
471 | | const AV1_COMP *cpi, MACROBLOCK *x, BLOCK_SIZE bsize, int dc_mode_cost, |
472 | | MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map, int64_t *best_rd, |
473 | | int *rate, int *rate_tokenonly, int64_t *distortion, int *skippable, |
474 | | int *beat_best_rd, PICK_MODE_CONTEXT *ctx, uint8_t *best_blk_skip, |
475 | 0 | uint8_t *tx_type_map) { |
476 | 0 | MACROBLOCKD *const xd = &x->e_mbd; |
477 | 0 | MB_MODE_INFO *const mbmi = xd->mi[0]; |
478 | 0 | assert(!is_inter_block(mbmi)); |
479 | 0 | assert(av1_allow_palette(cpi->common.features.allow_screen_content_tools, |
480 | 0 | bsize)); |
481 | 0 | assert(PALETTE_MAX_SIZE == 8); |
482 | 0 | assert(PALETTE_MIN_SIZE == 2); |
483 | |
|
484 | 0 | const int src_stride = x->plane[0].src.stride; |
485 | 0 | const uint8_t *const src = x->plane[0].src.buf; |
486 | 0 | int block_width, block_height, rows, cols; |
487 | 0 | av1_get_block_dimensions(bsize, 0, xd, &block_width, &block_height, &rows, |
488 | 0 | &cols); |
489 | 0 | const SequenceHeader *const seq_params = cpi->common.seq_params; |
490 | 0 | const int is_hbd = seq_params->use_highbitdepth; |
491 | 0 | const int bit_depth = seq_params->bit_depth; |
492 | 0 | int unused; |
493 | |
|
494 | 0 | int count_buf[1 << 12]; // Maximum (1 << 12) color levels. |
495 | 0 | int count_buf_8bit[1 << 8]; // Maximum (1 << 8) bins for hbd path. |
496 | 0 | int colors, colors_threshold = 0; |
497 | 0 | if (is_hbd) { |
498 | 0 | av1_count_colors_highbd(src, src_stride, rows, cols, bit_depth, count_buf, |
499 | 0 | count_buf_8bit, &colors_threshold, &colors); |
500 | 0 | } else { |
501 | 0 | av1_count_colors(src, src_stride, rows, cols, count_buf, &colors); |
502 | 0 | colors_threshold = colors; |
503 | 0 | } |
504 | |
|
505 | 0 | uint8_t *const color_map = xd->plane[0].color_index_map; |
506 | 0 | if (colors_threshold > 1 && colors_threshold <= 64) { |
507 | 0 | int *const data = x->palette_buffer->kmeans_data_buf; |
508 | 0 | int centroids[PALETTE_MAX_SIZE]; |
509 | 0 | int lower_bound, upper_bound; |
510 | 0 | fill_data_and_get_bounds(src, src_stride, rows, cols, is_hbd, data, |
511 | 0 | &lower_bound, &upper_bound); |
512 | |
|
513 | 0 | mbmi->mode = DC_PRED; |
514 | 0 | mbmi->filter_intra_mode_info.use_filter_intra = 0; |
515 | |
|
516 | 0 | uint16_t color_cache[2 * PALETTE_MAX_SIZE]; |
517 | 0 | const int n_cache = av1_get_palette_cache(xd, 0, color_cache); |
518 | | |
519 | | // Find the dominant colors, stored in top_colors[]. |
520 | 0 | int top_colors[PALETTE_MAX_SIZE] = { 0 }; |
521 | 0 | for (int i = 0; i < AOMMIN(colors, PALETTE_MAX_SIZE); ++i) { |
522 | 0 | int max_count = 0; |
523 | 0 | for (int j = 0; j < (1 << bit_depth); ++j) { |
524 | 0 | if (count_buf[j] > max_count) { |
525 | 0 | max_count = count_buf[j]; |
526 | 0 | top_colors[i] = j; |
527 | 0 | } |
528 | 0 | } |
529 | 0 | assert(max_count > 0); |
530 | 0 | count_buf[top_colors[i]] = 0; |
531 | 0 | } |
532 | | |
533 | | // The following are the approaches used for header rdcost based gating |
534 | | // for early termination for different values of prune_palette_search_level. |
535 | | // 0: Pruning based on header rdcost for ascending order palette_size |
536 | | // search. |
537 | | // 1: When colors > PALETTE_MIN_SIZE, enabled only for coarse palette_size |
538 | | // search and for finer search do_header_rd_based_gating parameter is |
539 | | // explicitly passed as 'false'. |
540 | | // 2: Enabled only for ascending order palette_size search and for |
541 | | // descending order search do_header_rd_based_gating parameter is explicitly |
542 | | // passed as 'false'. |
543 | 0 | const bool do_header_rd_based_gating = |
544 | 0 | cpi->sf.intra_sf.prune_luma_palette_size_search_level != 0; |
545 | | |
546 | | // TODO(huisu@google.com): Try to avoid duplicate computation in cases |
547 | | // where the dominant colors and the k-means results are similar. |
548 | 0 | if ((cpi->sf.intra_sf.prune_palette_search_level == 1) && |
549 | 0 | (colors > PALETTE_MIN_SIZE)) { |
550 | | // Start index and step size below are chosen to evaluate unique |
551 | | // candidates in neighbor search, in case a winner candidate is found in |
552 | | // coarse search. Example, |
553 | | // 1) 8 colors (end_n = 8): 2,3,4,5,6,7,8. start_n is chosen as 2 and step |
554 | | // size is chosen as 3. Therefore, coarse search will evaluate 2, 5 and 8. |
555 | | // If winner is found at 5, then 4 and 6 are evaluated. Similarly, for 2 |
556 | | // (3) and 8 (7). |
557 | | // 2) 7 colors (end_n = 7): 2,3,4,5,6,7. If start_n is chosen as 2 (same |
558 | | // as for 8 colors) then step size should also be 2, to cover all |
559 | | // candidates. Coarse search will evaluate 2, 4 and 6. If winner is either |
560 | | // 2 or 4, 3 will be evaluated. Instead, if start_n=3 and step_size=3, |
561 | | // coarse search will evaluate 3 and 6. For the winner, unique neighbors |
562 | | // (3: 2,4 or 6: 5,7) would be evaluated. |
563 | | |
564 | | // Start index for coarse palette search for dominant colors and k-means |
565 | 0 | const uint8_t start_n_lookup_table[PALETTE_MAX_SIZE + 1] = { 0, 0, 0, |
566 | 0 | 3, 3, 2, |
567 | 0 | 3, 3, 2 }; |
568 | | // Step size for coarse palette search for dominant colors and k-means |
569 | 0 | const uint8_t step_size_lookup_table[PALETTE_MAX_SIZE + 1] = { 0, 0, 0, |
570 | 0 | 3, 3, 3, |
571 | 0 | 3, 3, 3 }; |
572 | | |
573 | | // Choose the start index and step size for coarse search based on number |
574 | | // of colors |
575 | 0 | const int max_n = AOMMIN(colors, PALETTE_MAX_SIZE); |
576 | 0 | const int min_n = start_n_lookup_table[max_n]; |
577 | 0 | const int step_size = step_size_lookup_table[max_n]; |
578 | 0 | assert(min_n >= PALETTE_MIN_SIZE); |
579 | | // Perform top color coarse palette search to find the winner candidate |
580 | 0 | const int top_color_winner = perform_top_color_palette_search( |
581 | 0 | cpi, x, mbmi, bsize, dc_mode_cost, data, top_colors, min_n, max_n + 1, |
582 | 0 | step_size, do_header_rd_based_gating, &unused, color_cache, n_cache, |
583 | 0 | best_mbmi, best_palette_color_map, best_rd, rate, rate_tokenonly, |
584 | 0 | distortion, skippable, beat_best_rd, ctx, best_blk_skip, tx_type_map); |
585 | | // Evaluate neighbors for the winner color (if winner is found) in the |
586 | | // above coarse search for dominant colors |
587 | 0 | if (top_color_winner <= max_n) { |
588 | 0 | int stage2_min_n, stage2_max_n, stage2_step_size; |
589 | 0 | set_stage2_params(&stage2_min_n, &stage2_max_n, &stage2_step_size, |
590 | 0 | top_color_winner, max_n); |
591 | | // perform finer search for the winner candidate |
592 | 0 | perform_top_color_palette_search( |
593 | 0 | cpi, x, mbmi, bsize, dc_mode_cost, data, top_colors, stage2_min_n, |
594 | 0 | stage2_max_n + 1, stage2_step_size, |
595 | | /*do_header_rd_based_gating=*/false, &unused, color_cache, n_cache, |
596 | 0 | best_mbmi, best_palette_color_map, best_rd, rate, rate_tokenonly, |
597 | 0 | distortion, skippable, beat_best_rd, ctx, best_blk_skip, |
598 | 0 | tx_type_map); |
599 | 0 | } |
600 | | // K-means clustering. |
601 | | // Perform k-means coarse palette search to find the winner candidate |
602 | 0 | const int k_means_winner = perform_k_means_palette_search( |
603 | 0 | cpi, x, mbmi, bsize, dc_mode_cost, data, lower_bound, upper_bound, |
604 | 0 | min_n, max_n + 1, step_size, do_header_rd_based_gating, &unused, |
605 | 0 | color_cache, n_cache, best_mbmi, best_palette_color_map, best_rd, |
606 | 0 | rate, rate_tokenonly, distortion, skippable, beat_best_rd, ctx, |
607 | 0 | best_blk_skip, tx_type_map, color_map, rows * cols); |
608 | | // Evaluate neighbors for the winner color (if winner is found) in the |
609 | | // above coarse search for k-means |
610 | 0 | if (k_means_winner <= max_n) { |
611 | 0 | int start_n_stage2, end_n_stage2, step_size_stage2; |
612 | 0 | set_stage2_params(&start_n_stage2, &end_n_stage2, &step_size_stage2, |
613 | 0 | k_means_winner, max_n); |
614 | | // perform finer search for the winner candidate |
615 | 0 | perform_k_means_palette_search( |
616 | 0 | cpi, x, mbmi, bsize, dc_mode_cost, data, lower_bound, upper_bound, |
617 | 0 | start_n_stage2, end_n_stage2 + 1, step_size_stage2, |
618 | | /*do_header_rd_based_gating=*/false, &unused, color_cache, n_cache, |
619 | 0 | best_mbmi, best_palette_color_map, best_rd, rate, rate_tokenonly, |
620 | 0 | distortion, skippable, beat_best_rd, ctx, best_blk_skip, |
621 | 0 | tx_type_map, color_map, rows * cols); |
622 | 0 | } |
623 | 0 | } else { |
624 | 0 | const int max_n = AOMMIN(colors, PALETTE_MAX_SIZE), |
625 | 0 | min_n = PALETTE_MIN_SIZE; |
626 | | // Perform top color palette search in ascending order |
627 | 0 | int last_n_searched = min_n; |
628 | 0 | perform_top_color_palette_search( |
629 | 0 | cpi, x, mbmi, bsize, dc_mode_cost, data, top_colors, min_n, max_n + 1, |
630 | 0 | 1, do_header_rd_based_gating, &last_n_searched, color_cache, n_cache, |
631 | 0 | best_mbmi, best_palette_color_map, best_rd, rate, rate_tokenonly, |
632 | 0 | distortion, skippable, beat_best_rd, ctx, best_blk_skip, tx_type_map); |
633 | 0 | if (last_n_searched < max_n) { |
634 | | // Search in descending order until we get to the previous best |
635 | 0 | perform_top_color_palette_search( |
636 | 0 | cpi, x, mbmi, bsize, dc_mode_cost, data, top_colors, max_n, |
637 | 0 | last_n_searched, -1, /*do_header_rd_based_gating=*/false, &unused, |
638 | 0 | color_cache, n_cache, best_mbmi, best_palette_color_map, best_rd, |
639 | 0 | rate, rate_tokenonly, distortion, skippable, beat_best_rd, ctx, |
640 | 0 | best_blk_skip, tx_type_map); |
641 | 0 | } |
642 | | // K-means clustering. |
643 | 0 | if (colors == PALETTE_MIN_SIZE) { |
644 | | // Special case: These colors automatically become the centroids. |
645 | 0 | assert(colors == 2); |
646 | 0 | centroids[0] = lower_bound; |
647 | 0 | centroids[1] = upper_bound; |
648 | 0 | palette_rd_y(cpi, x, mbmi, bsize, dc_mode_cost, data, centroids, colors, |
649 | 0 | color_cache, n_cache, /*do_header_rd_based_gating=*/false, |
650 | 0 | best_mbmi, best_palette_color_map, best_rd, rate, |
651 | 0 | rate_tokenonly, distortion, skippable, beat_best_rd, ctx, |
652 | 0 | best_blk_skip, tx_type_map, NULL, NULL); |
653 | 0 | } else { |
654 | | // Perform k-means palette search in ascending order |
655 | 0 | last_n_searched = min_n; |
656 | 0 | perform_k_means_palette_search( |
657 | 0 | cpi, x, mbmi, bsize, dc_mode_cost, data, lower_bound, upper_bound, |
658 | 0 | min_n, max_n + 1, 1, do_header_rd_based_gating, &last_n_searched, |
659 | 0 | color_cache, n_cache, best_mbmi, best_palette_color_map, best_rd, |
660 | 0 | rate, rate_tokenonly, distortion, skippable, beat_best_rd, ctx, |
661 | 0 | best_blk_skip, tx_type_map, color_map, rows * cols); |
662 | 0 | if (last_n_searched < max_n) { |
663 | | // Search in descending order until we get to the previous best |
664 | 0 | perform_k_means_palette_search( |
665 | 0 | cpi, x, mbmi, bsize, dc_mode_cost, data, lower_bound, upper_bound, |
666 | 0 | max_n, last_n_searched, -1, /*do_header_rd_based_gating=*/false, |
667 | 0 | &unused, color_cache, n_cache, best_mbmi, best_palette_color_map, |
668 | 0 | best_rd, rate, rate_tokenonly, distortion, skippable, |
669 | 0 | beat_best_rd, ctx, best_blk_skip, tx_type_map, color_map, |
670 | 0 | rows * cols); |
671 | 0 | } |
672 | 0 | } |
673 | 0 | } |
674 | 0 | } |
675 | |
|
676 | 0 | if (best_mbmi->palette_mode_info.palette_size[0] > 0) { |
677 | 0 | memcpy(color_map, best_palette_color_map, |
678 | 0 | block_width * block_height * sizeof(best_palette_color_map[0])); |
679 | 0 | } |
680 | 0 | *mbmi = *best_mbmi; |
681 | 0 | } |
682 | | |
683 | | void av1_rd_pick_palette_intra_sbuv(const AV1_COMP *cpi, MACROBLOCK *x, |
684 | | int dc_mode_cost, |
685 | | uint8_t *best_palette_color_map, |
686 | | MB_MODE_INFO *const best_mbmi, |
687 | | int64_t *best_rd, int *rate, |
688 | | int *rate_tokenonly, int64_t *distortion, |
689 | 0 | int *skippable) { |
690 | 0 | MACROBLOCKD *const xd = &x->e_mbd; |
691 | 0 | MB_MODE_INFO *const mbmi = xd->mi[0]; |
692 | 0 | assert(!is_inter_block(mbmi)); |
693 | 0 | assert(av1_allow_palette(cpi->common.features.allow_screen_content_tools, |
694 | 0 | mbmi->bsize)); |
695 | 0 | PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info; |
696 | 0 | const BLOCK_SIZE bsize = mbmi->bsize; |
697 | 0 | const SequenceHeader *const seq_params = cpi->common.seq_params; |
698 | 0 | int this_rate; |
699 | 0 | int64_t this_rd; |
700 | 0 | int colors_u, colors_v; |
701 | 0 | int colors_threshold_u = 0, colors_threshold_v = 0, colors_threshold = 0; |
702 | 0 | const int src_stride = x->plane[1].src.stride; |
703 | 0 | const uint8_t *const src_u = x->plane[1].src.buf; |
704 | 0 | const uint8_t *const src_v = x->plane[2].src.buf; |
705 | 0 | uint8_t *const color_map = xd->plane[1].color_index_map; |
706 | 0 | RD_STATS tokenonly_rd_stats; |
707 | 0 | int plane_block_width, plane_block_height, rows, cols; |
708 | 0 | av1_get_block_dimensions(bsize, 1, xd, &plane_block_width, |
709 | 0 | &plane_block_height, &rows, &cols); |
710 | |
|
711 | 0 | mbmi->uv_mode = UV_DC_PRED; |
712 | 0 | int count_buf[1 << 12]; // Maximum (1 << 12) color levels. |
713 | 0 | int count_buf_8bit[1 << 8]; // Maximum (1 << 8) bins for hbd path. |
714 | 0 | if (seq_params->use_highbitdepth) { |
715 | 0 | av1_count_colors_highbd(src_u, src_stride, rows, cols, |
716 | 0 | seq_params->bit_depth, count_buf, count_buf_8bit, |
717 | 0 | &colors_threshold_u, &colors_u); |
718 | 0 | av1_count_colors_highbd(src_v, src_stride, rows, cols, |
719 | 0 | seq_params->bit_depth, count_buf, count_buf_8bit, |
720 | 0 | &colors_threshold_v, &colors_v); |
721 | 0 | } else { |
722 | 0 | av1_count_colors(src_u, src_stride, rows, cols, count_buf, &colors_u); |
723 | 0 | av1_count_colors(src_v, src_stride, rows, cols, count_buf, &colors_v); |
724 | 0 | colors_threshold_u = colors_u; |
725 | 0 | colors_threshold_v = colors_v; |
726 | 0 | } |
727 | |
|
728 | 0 | uint16_t color_cache[2 * PALETTE_MAX_SIZE]; |
729 | 0 | const int n_cache = av1_get_palette_cache(xd, 1, color_cache); |
730 | |
|
731 | 0 | colors_threshold = colors_threshold_u > colors_threshold_v |
732 | 0 | ? colors_threshold_u |
733 | 0 | : colors_threshold_v; |
734 | 0 | if (colors_threshold > 1 && colors_threshold <= 64) { |
735 | 0 | int r, c, n, i, j; |
736 | 0 | const int max_itr = 50; |
737 | 0 | int lb_u, ub_u, val_u; |
738 | 0 | int lb_v, ub_v, val_v; |
739 | 0 | int *const data = x->palette_buffer->kmeans_data_buf; |
740 | 0 | int centroids[2 * PALETTE_MAX_SIZE]; |
741 | |
|
742 | 0 | uint16_t *src_u16 = CONVERT_TO_SHORTPTR(src_u); |
743 | 0 | uint16_t *src_v16 = CONVERT_TO_SHORTPTR(src_v); |
744 | 0 | if (seq_params->use_highbitdepth) { |
745 | 0 | lb_u = src_u16[0]; |
746 | 0 | ub_u = src_u16[0]; |
747 | 0 | lb_v = src_v16[0]; |
748 | 0 | ub_v = src_v16[0]; |
749 | 0 | } else { |
750 | 0 | lb_u = src_u[0]; |
751 | 0 | ub_u = src_u[0]; |
752 | 0 | lb_v = src_v[0]; |
753 | 0 | ub_v = src_v[0]; |
754 | 0 | } |
755 | |
|
756 | 0 | for (r = 0; r < rows; ++r) { |
757 | 0 | for (c = 0; c < cols; ++c) { |
758 | 0 | if (seq_params->use_highbitdepth) { |
759 | 0 | val_u = src_u16[r * src_stride + c]; |
760 | 0 | val_v = src_v16[r * src_stride + c]; |
761 | 0 | data[(r * cols + c) * 2] = val_u; |
762 | 0 | data[(r * cols + c) * 2 + 1] = val_v; |
763 | 0 | } else { |
764 | 0 | val_u = src_u[r * src_stride + c]; |
765 | 0 | val_v = src_v[r * src_stride + c]; |
766 | 0 | data[(r * cols + c) * 2] = val_u; |
767 | 0 | data[(r * cols + c) * 2 + 1] = val_v; |
768 | 0 | } |
769 | 0 | if (val_u < lb_u) |
770 | 0 | lb_u = val_u; |
771 | 0 | else if (val_u > ub_u) |
772 | 0 | ub_u = val_u; |
773 | 0 | if (val_v < lb_v) |
774 | 0 | lb_v = val_v; |
775 | 0 | else if (val_v > ub_v) |
776 | 0 | ub_v = val_v; |
777 | 0 | } |
778 | 0 | } |
779 | |
|
780 | 0 | const int colors = colors_u > colors_v ? colors_u : colors_v; |
781 | 0 | const int max_colors = |
782 | 0 | colors > PALETTE_MAX_SIZE ? PALETTE_MAX_SIZE : colors; |
783 | 0 | for (n = PALETTE_MIN_SIZE; n <= max_colors; ++n) { |
784 | 0 | for (i = 0; i < n; ++i) { |
785 | 0 | centroids[i * 2] = lb_u + (2 * i + 1) * (ub_u - lb_u) / n / 2; |
786 | 0 | centroids[i * 2 + 1] = lb_v + (2 * i + 1) * (ub_v - lb_v) / n / 2; |
787 | 0 | } |
788 | 0 | av1_k_means(data, centroids, color_map, rows * cols, n, 2, max_itr); |
789 | 0 | optimize_palette_colors(color_cache, n_cache, n, 2, centroids, |
790 | 0 | cpi->common.seq_params->bit_depth); |
791 | | // Sort the U channel colors in ascending order. |
792 | 0 | for (i = 0; i < 2 * (n - 1); i += 2) { |
793 | 0 | int min_idx = i; |
794 | 0 | int min_val = centroids[i]; |
795 | 0 | for (j = i + 2; j < 2 * n; j += 2) |
796 | 0 | if (centroids[j] < min_val) min_val = centroids[j], min_idx = j; |
797 | 0 | if (min_idx != i) { |
798 | 0 | int temp_u = centroids[i], temp_v = centroids[i + 1]; |
799 | 0 | centroids[i] = centroids[min_idx]; |
800 | 0 | centroids[i + 1] = centroids[min_idx + 1]; |
801 | 0 | centroids[min_idx] = temp_u, centroids[min_idx + 1] = temp_v; |
802 | 0 | } |
803 | 0 | } |
804 | 0 | av1_calc_indices(data, centroids, color_map, rows * cols, n, 2); |
805 | 0 | extend_palette_color_map(color_map, cols, rows, plane_block_width, |
806 | 0 | plane_block_height); |
807 | 0 | pmi->palette_size[1] = n; |
808 | 0 | for (i = 1; i < 3; ++i) { |
809 | 0 | for (j = 0; j < n; ++j) { |
810 | 0 | if (seq_params->use_highbitdepth) |
811 | 0 | pmi->palette_colors[i * PALETTE_MAX_SIZE + j] = clip_pixel_highbd( |
812 | 0 | (int)centroids[j * 2 + i - 1], seq_params->bit_depth); |
813 | 0 | else |
814 | 0 | pmi->palette_colors[i * PALETTE_MAX_SIZE + j] = |
815 | 0 | clip_pixel((int)centroids[j * 2 + i - 1]); |
816 | 0 | } |
817 | 0 | } |
818 | |
|
819 | 0 | if (cpi->sf.intra_sf.early_term_chroma_palette_size_search) { |
820 | 0 | const int palette_mode_rate = |
821 | 0 | intra_mode_info_cost_uv(cpi, x, mbmi, bsize, dc_mode_cost); |
822 | 0 | const int64_t header_rd = RDCOST(x->rdmult, palette_mode_rate, 0); |
823 | | // Terminate further palette_size search, if header cost corresponding |
824 | | // to lower palette_size is more than the best_rd. |
825 | 0 | if (header_rd >= *best_rd) break; |
826 | 0 | av1_txfm_uvrd(cpi, x, &tokenonly_rd_stats, bsize, *best_rd); |
827 | 0 | if (tokenonly_rd_stats.rate == INT_MAX) continue; |
828 | 0 | this_rate = tokenonly_rd_stats.rate + palette_mode_rate; |
829 | 0 | } else { |
830 | 0 | av1_txfm_uvrd(cpi, x, &tokenonly_rd_stats, bsize, *best_rd); |
831 | 0 | if (tokenonly_rd_stats.rate == INT_MAX) continue; |
832 | 0 | this_rate = tokenonly_rd_stats.rate + |
833 | 0 | intra_mode_info_cost_uv(cpi, x, mbmi, bsize, dc_mode_cost); |
834 | 0 | } |
835 | | |
836 | 0 | this_rd = RDCOST(x->rdmult, this_rate, tokenonly_rd_stats.dist); |
837 | 0 | if (this_rd < *best_rd) { |
838 | 0 | *best_rd = this_rd; |
839 | 0 | *best_mbmi = *mbmi; |
840 | 0 | memcpy(best_palette_color_map, color_map, |
841 | 0 | plane_block_width * plane_block_height * |
842 | 0 | sizeof(best_palette_color_map[0])); |
843 | 0 | *rate = this_rate; |
844 | 0 | *distortion = tokenonly_rd_stats.dist; |
845 | 0 | *rate_tokenonly = tokenonly_rd_stats.rate; |
846 | 0 | *skippable = tokenonly_rd_stats.skip_txfm; |
847 | 0 | } |
848 | 0 | } |
849 | 0 | } |
850 | 0 | if (best_mbmi->palette_mode_info.palette_size[1] > 0) { |
851 | 0 | memcpy(color_map, best_palette_color_map, |
852 | 0 | plane_block_width * plane_block_height * |
853 | 0 | sizeof(best_palette_color_map[0])); |
854 | 0 | } |
855 | 0 | } |
856 | | |
857 | 0 | void av1_restore_uv_color_map(const AV1_COMP *cpi, MACROBLOCK *x) { |
858 | 0 | MACROBLOCKD *const xd = &x->e_mbd; |
859 | 0 | MB_MODE_INFO *const mbmi = xd->mi[0]; |
860 | 0 | PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info; |
861 | 0 | const BLOCK_SIZE bsize = mbmi->bsize; |
862 | 0 | int src_stride = x->plane[1].src.stride; |
863 | 0 | const uint8_t *const src_u = x->plane[1].src.buf; |
864 | 0 | const uint8_t *const src_v = x->plane[2].src.buf; |
865 | 0 | int *const data = x->palette_buffer->kmeans_data_buf; |
866 | 0 | int centroids[2 * PALETTE_MAX_SIZE]; |
867 | 0 | uint8_t *const color_map = xd->plane[1].color_index_map; |
868 | 0 | int r, c; |
869 | 0 | const uint16_t *const src_u16 = CONVERT_TO_SHORTPTR(src_u); |
870 | 0 | const uint16_t *const src_v16 = CONVERT_TO_SHORTPTR(src_v); |
871 | 0 | int plane_block_width, plane_block_height, rows, cols; |
872 | 0 | av1_get_block_dimensions(bsize, 1, xd, &plane_block_width, |
873 | 0 | &plane_block_height, &rows, &cols); |
874 | |
|
875 | 0 | for (r = 0; r < rows; ++r) { |
876 | 0 | for (c = 0; c < cols; ++c) { |
877 | 0 | if (cpi->common.seq_params->use_highbitdepth) { |
878 | 0 | data[(r * cols + c) * 2] = src_u16[r * src_stride + c]; |
879 | 0 | data[(r * cols + c) * 2 + 1] = src_v16[r * src_stride + c]; |
880 | 0 | } else { |
881 | 0 | data[(r * cols + c) * 2] = src_u[r * src_stride + c]; |
882 | 0 | data[(r * cols + c) * 2 + 1] = src_v[r * src_stride + c]; |
883 | 0 | } |
884 | 0 | } |
885 | 0 | } |
886 | |
|
887 | 0 | for (r = 1; r < 3; ++r) { |
888 | 0 | for (c = 0; c < pmi->palette_size[1]; ++c) { |
889 | 0 | centroids[c * 2 + r - 1] = pmi->palette_colors[r * PALETTE_MAX_SIZE + c]; |
890 | 0 | } |
891 | 0 | } |
892 | |
|
893 | 0 | av1_calc_indices(data, centroids, color_map, rows * cols, |
894 | 0 | pmi->palette_size[1], 2); |
895 | 0 | extend_palette_color_map(color_map, cols, rows, plane_block_width, |
896 | 0 | plane_block_height); |
897 | 0 | } |