av1/encoder/tokenize.c - aom - Git at Google

 /*
  * Copyright (c) 2016, Alliance for Open Media. All rights reserved
  *
  * This source code is subject to the terms of the BSD 2 Clause License and
  * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
  * was not distributed with this source code in the LICENSE file, you can
  * obtain it at www.aomedia.org/license/software. 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 www.aomedia.org/license/patent.
  */

 #include <assert.h>
 #include <math.h>
 #include <stdio.h>
 #include <string.h>

 #include "aom_mem/aom_mem.h"

 #include "av1/common/entropy.h"
 #include "av1/common/pred_common.h"
 #include "av1/common/scan.h"
 #include "av1/common/seg_common.h"

 #include "av1/encoder/cost.h"
 #include "av1/encoder/encoder.h"
 #include "av1/encoder/encodetxb.h"
 #include "av1/encoder/rdopt.h"
 #include "av1/encoder/tokenize.h"

 static int cost_and_tokenize_map(Av1ColorMapParam *param, TokenExtra **t,
                                  int plane, int calc_rate, int allow_update_cdf,
                                  FRAME_COUNTS *counts) {
   const uint8_t *const color_map = param->color_map;
   MapCdf map_cdf = param->map_cdf;
   ColorCost color_cost = param->color_cost;
   const int plane_block_width = param->plane_width;
   const int rows = param->rows;
   const int cols = param->cols;
   const int n = param->n_colors;
   const int palette_size_idx = n - PALETTE_MIN_SIZE;
   int this_rate = 0;

   (void)plane;
   (void)counts;

   for (int k = 1; k < rows + cols - 1; ++k) {
     for (int j = AOMMIN(k, cols - 1); j >= AOMMAX(0, k - rows + 1); --j) {
       int i = k - j;
       int color_new_idx;
       const int color_ctx = av1_fast_palette_color_index_context(
           color_map, plane_block_width, i, j, &color_new_idx);
       assert(color_new_idx >= 0 && color_new_idx < n);
       if (calc_rate) {
         this_rate += color_cost[palette_size_idx][color_ctx][color_new_idx];
       } else {
         (*t)->token = color_new_idx;
         (*t)->color_ctx = color_ctx;
         ++(*t);
         if (allow_update_cdf)
           update_cdf(map_cdf[palette_size_idx][color_ctx], color_new_idx, n);
 #if CONFIG_ENTROPY_STATS
         if (plane) {
           ++counts->palette_uv_color_index[palette_size_idx][color_ctx]
                                           [color_new_idx];
         } else {
           ++counts->palette_y_color_index[palette_size_idx][color_ctx]
                                          [color_new_idx];
         }
 #endif
       }
     }
   }
   if (calc_rate) return this_rate;
   return 0;
 }

 static void get_palette_params(const MACROBLOCK *const x, int plane,
                                BLOCK_SIZE bsize, Av1ColorMapParam *params) {
   const MACROBLOCKD *const xd = &x->e_mbd;
   const MB_MODE_INFO *const mbmi = xd->mi[0];
   const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
   params->color_map = xd->plane[plane].color_index_map;
   params->map_cdf = plane ? xd->tile_ctx->palette_uv_color_index_cdf
                           : xd->tile_ctx->palette_y_color_index_cdf;
   params->color_cost = plane ? x->mode_costs.palette_uv_color_cost
                              : x->mode_costs.palette_y_color_cost;
   params->n_colors = pmi->palette_size[plane];
   av1_get_block_dimensions(bsize, plane, xd, &params->plane_width, NULL,
                            &params->rows, &params->cols);
 }

 // TODO(any): Remove this function
 static void get_color_map_params(const MACROBLOCK *const x, int plane,
                                  BLOCK_SIZE bsize, TX_SIZE tx_size,
                                  COLOR_MAP_TYPE type,
                                  Av1ColorMapParam *params) {
   (void)tx_size;
   memset(params, 0, sizeof(*params));
   switch (type) {
     case PALETTE_MAP: get_palette_params(x, plane, bsize, params); break;
     default: assert(0 && "Invalid color map type"); return;
   }
 }

 int av1_cost_color_map(const MACROBLOCK *const x, int plane, BLOCK_SIZE bsize,
                        TX_SIZE tx_size, COLOR_MAP_TYPE type) {
   assert(plane == 0 || plane == 1);
   Av1ColorMapParam color_map_params;
   get_color_map_params(x, plane, bsize, tx_size, type, &color_map_params);
   return cost_and_tokenize_map(&color_map_params, NULL, plane, 1, 0, NULL);
 }

 void av1_tokenize_color_map(const MACROBLOCK *const x, int plane,
                             TokenExtra **t, BLOCK_SIZE bsize, TX_SIZE tx_size,
                             COLOR_MAP_TYPE type, int allow_update_cdf,
                             FRAME_COUNTS *counts) {
   assert(plane == 0 || plane == 1);
   Av1ColorMapParam color_map_params;
   get_color_map_params(x, plane, bsize, tx_size, type, &color_map_params);
   // The first color index does not use context or entropy.
   (*t)->token = color_map_params.color_map[0];
   (*t)->color_ctx = -1;
   ++(*t);
   cost_and_tokenize_map(&color_map_params, t, plane, 0, allow_update_cdf,
                         counts);
 }

 static void tokenize_vartx(ThreadData *td, TX_SIZE tx_size,
                            BLOCK_SIZE plane_bsize, int blk_row, int blk_col,
                            int block, int plane, void *arg) {
   MACROBLOCK *const x = &td->mb;
   MACROBLOCKD *const xd = &x->e_mbd;
   MB_MODE_INFO *const mbmi = xd->mi[0];
   const struct macroblockd_plane *const pd = &xd->plane[plane];
   const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
   const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);

   if (blk_row >= max_blocks_high || blk_col >= max_blocks_wide) return;

   const TX_SIZE plane_tx_size =
       plane ? av1_get_max_uv_txsize(mbmi->bsize, pd->subsampling_x,
                                     pd->subsampling_y)
             : mbmi->inter_tx_size[av1_get_txb_size_index(plane_bsize, blk_row,
                                                          blk_col)];

   if (tx_size == plane_tx_size || plane) {
     plane_bsize =
         get_plane_block_size(mbmi->bsize, pd->subsampling_x, pd->subsampling_y);
     av1_update_and_record_txb_context(plane, block, blk_row, blk_col,
                                       plane_bsize, tx_size, arg);

   } else {
     // Half the block size in transform block unit.
     const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
     const int bsw = tx_size_wide_unit[sub_txs];
     const int bsh = tx_size_high_unit[sub_txs];
     const int step = bsw * bsh;
     const int row_end =
         AOMMIN(tx_size_high_unit[tx_size], max_blocks_high - blk_row);
     const int col_end =
         AOMMIN(tx_size_wide_unit[tx_size], max_blocks_wide - blk_col);

     assert(bsw > 0 && bsh > 0);

     for (int row = 0; row < row_end; row += bsh) {
       const int offsetr = blk_row + row;
       for (int col = 0; col < col_end; col += bsw) {
         const int offsetc = blk_col + col;

         tokenize_vartx(td, sub_txs, plane_bsize, offsetr, offsetc, block, plane,
                        arg);
         block += step;
       }
     }
   }
 }

 void av1_tokenize_sb_vartx(const AV1_COMP *cpi, ThreadData *td,
                            RUN_TYPE dry_run, BLOCK_SIZE bsize, int *rate,
                            uint8_t allow_update_cdf) {
   assert(bsize < BLOCK_SIZES_ALL);
   const AV1_COMMON *const cm = &cpi->common;
   MACROBLOCK *const x = &td->mb;
   MACROBLOCKD *const xd = &x->e_mbd;
   const int mi_row = xd->mi_row;
   const int mi_col = xd->mi_col;
   if (mi_row >= cm->mi_params.mi_rows || mi_col >= cm->mi_params.mi_cols)
     return;

   const int num_planes = av1_num_planes(cm);
   MB_MODE_INFO *const mbmi = xd->mi[0];
   struct tokenize_b_args arg = { cpi, td, 0, allow_update_cdf, dry_run };

   if (mbmi->skip_txfm) {
     av1_reset_entropy_context(xd, bsize, num_planes);
     return;
   }

   for (int plane = 0; plane < num_planes; ++plane) {
     if (plane && !xd->is_chroma_ref) break;
     const struct macroblockd_plane *const pd = &xd->plane[plane];
     const int ss_x = pd->subsampling_x;
     const int ss_y = pd->subsampling_y;
     const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, ss_x, ss_y);
     assert(plane_bsize < BLOCK_SIZES_ALL);
     const int mi_width = mi_size_wide[plane_bsize];
     const int mi_height = mi_size_high[plane_bsize];
     const TX_SIZE max_tx_size = get_vartx_max_txsize(xd, plane_bsize, plane);
     const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
     const int bw = mi_size_wide[txb_size];
     const int bh = mi_size_high[txb_size];
     int block = 0;
     const int step =
         tx_size_wide_unit[max_tx_size] * tx_size_high_unit[max_tx_size];

     const BLOCK_SIZE max_unit_bsize =
         get_plane_block_size(BLOCK_64X64, ss_x, ss_y);
     int mu_blocks_wide = mi_size_wide[max_unit_bsize];
     int mu_blocks_high = mi_size_high[max_unit_bsize];

     mu_blocks_wide = AOMMIN(mi_width, mu_blocks_wide);
     mu_blocks_high = AOMMIN(mi_height, mu_blocks_high);

     for (int idy = 0; idy < mi_height; idy += mu_blocks_high) {
       for (int idx = 0; idx < mi_width; idx += mu_blocks_wide) {
         const int unit_height = AOMMIN(mu_blocks_high + idy, mi_height);
         const int unit_width = AOMMIN(mu_blocks_wide + idx, mi_width);
         for (int blk_row = idy; blk_row < unit_height; blk_row += bh) {
           for (int blk_col = idx; blk_col < unit_width; blk_col += bw) {
             tokenize_vartx(td, max_tx_size, plane_bsize, blk_row, blk_col,
                            block, plane, &arg);
             block += step;
           }
         }
       }
     }
   }
   if (rate) *rate += arg.this_rate;
 }
	/*
	* Copyright (c) 2016, Alliance for Open Media. All rights reserved
	*
	* This source code is subject to the terms of the BSD 2 Clause License and
	* the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
	* was not distributed with this source code in the LICENSE file, you can
	* obtain it at www.aomedia.org/license/software. 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 www.aomedia.org/license/patent.
	*/

	#include <assert.h>
	#include <math.h>
	#include <stdio.h>
	#include <string.h>

	#include "aom_mem/aom_mem.h"

	#include "av1/common/entropy.h"
	#include "av1/common/pred_common.h"
	#include "av1/common/scan.h"
	#include "av1/common/seg_common.h"

	#include "av1/encoder/cost.h"
	#include "av1/encoder/encoder.h"
	#include "av1/encoder/encodetxb.h"
	#include "av1/encoder/rdopt.h"
	#include "av1/encoder/tokenize.h"

	static int cost_and_tokenize_map(Av1ColorMapParam param, TokenExtra *t,
	int plane, int calc_rate, int allow_update_cdf,
	FRAME_COUNTS *counts) {
	const uint8_t *const color_map = param->color_map;
	MapCdf map_cdf = param->map_cdf;
	ColorCost color_cost = param->color_cost;
	const int plane_block_width = param->plane_width;
	const int rows = param->rows;
	const int cols = param->cols;
	const int n = param->n_colors;
	const int palette_size_idx = n - PALETTE_MIN_SIZE;
	int this_rate = 0;

	(void)plane;
	(void)counts;

	for (int k = 1; k < rows + cols - 1; ++k) {
	for (int j = AOMMIN(k, cols - 1); j >= AOMMAX(0, k - rows + 1); --j) {
	int i = k - j;
	int color_new_idx;
	const int color_ctx = av1_fast_palette_color_index_context(
	color_map, plane_block_width, i, j, &color_new_idx);
	assert(color_new_idx >= 0 && color_new_idx < n);
	if (calc_rate) {
	this_rate += color_cost[palette_size_idx][color_ctx][color_new_idx];
	} else {
	(*t)->token = color_new_idx;
	(*t)->color_ctx = color_ctx;
	++(*t);
	if (allow_update_cdf)
	update_cdf(map_cdf[palette_size_idx][color_ctx], color_new_idx, n);
	#if CONFIG_ENTROPY_STATS
	if (plane) {
	++counts->palette_uv_color_index[palette_size_idx][color_ctx]
	[color_new_idx];
	} else {
	++counts->palette_y_color_index[palette_size_idx][color_ctx]
	[color_new_idx];
	}
	#endif
	}
	}
	}
	if (calc_rate) return this_rate;
	return 0;
	}

	static void get_palette_params(const MACROBLOCK *const x, int plane,
	BLOCK_SIZE bsize, Av1ColorMapParam *params) {
	const MACROBLOCKD *const xd = &x->e_mbd;
	const MB_MODE_INFO *const mbmi = xd->mi[0];
	const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
	params->color_map = xd->plane[plane].color_index_map;
	params->map_cdf = plane ? xd->tile_ctx->palette_uv_color_index_cdf
	: xd->tile_ctx->palette_y_color_index_cdf;
	params->color_cost = plane ? x->mode_costs.palette_uv_color_cost
	: x->mode_costs.palette_y_color_cost;
	params->n_colors = pmi->palette_size[plane];
	av1_get_block_dimensions(bsize, plane, xd, &params->plane_width, NULL,
	&params->rows, &params->cols);
	}

	// TODO(any): Remove this function
	static void get_color_map_params(const MACROBLOCK *const x, int plane,
	BLOCK_SIZE bsize, TX_SIZE tx_size,
	COLOR_MAP_TYPE type,
	Av1ColorMapParam *params) {
	(void)tx_size;
	memset(params, 0, sizeof(*params));
	switch (type) {
	case PALETTE_MAP: get_palette_params(x, plane, bsize, params); break;
	default: assert(0 && "Invalid color map type"); return;
	}
	}

	int av1_cost_color_map(const MACROBLOCK *const x, int plane, BLOCK_SIZE bsize,
	TX_SIZE tx_size, COLOR_MAP_TYPE type) {
	assert(plane == 0 \|\| plane == 1);
	Av1ColorMapParam color_map_params;
	get_color_map_params(x, plane, bsize, tx_size, type, &color_map_params);
	return cost_and_tokenize_map(&color_map_params, NULL, plane, 1, 0, NULL);
	}

	void av1_tokenize_color_map(const MACROBLOCK *const x, int plane,
	TokenExtra **t, BLOCK_SIZE bsize, TX_SIZE tx_size,
	COLOR_MAP_TYPE type, int allow_update_cdf,
	FRAME_COUNTS *counts) {
	assert(plane == 0 \|\| plane == 1);
	Av1ColorMapParam color_map_params;
	get_color_map_params(x, plane, bsize, tx_size, type, &color_map_params);
	// The first color index does not use context or entropy.
	(*t)->token = color_map_params.color_map[0];
	(*t)->color_ctx = -1;
	++(*t);
	cost_and_tokenize_map(&color_map_params, t, plane, 0, allow_update_cdf,
	counts);
	}

	static void tokenize_vartx(ThreadData *td, TX_SIZE tx_size,
	BLOCK_SIZE plane_bsize, int blk_row, int blk_col,
	int block, int plane, void *arg) {
	MACROBLOCK *const x = &td->mb;
	MACROBLOCKD *const xd = &x->e_mbd;
	MB_MODE_INFO *const mbmi = xd->mi[0];
	const struct macroblockd_plane *const pd = &xd->plane[plane];
	const int max_blocks_high = max_block_high(xd, plane_bsize, plane);
	const int max_blocks_wide = max_block_wide(xd, plane_bsize, plane);

	if (blk_row >= max_blocks_high \|\| blk_col >= max_blocks_wide) return;

	const TX_SIZE plane_tx_size =
	plane ? av1_get_max_uv_txsize(mbmi->bsize, pd->subsampling_x,
	pd->subsampling_y)
	: mbmi->inter_tx_size[av1_get_txb_size_index(plane_bsize, blk_row,
	blk_col)];

	if (tx_size == plane_tx_size \|\| plane) {
	plane_bsize =
	get_plane_block_size(mbmi->bsize, pd->subsampling_x, pd->subsampling_y);
	av1_update_and_record_txb_context(plane, block, blk_row, blk_col,
	plane_bsize, tx_size, arg);

	} else {
	// Half the block size in transform block unit.
	const TX_SIZE sub_txs = sub_tx_size_map[tx_size];
	const int bsw = tx_size_wide_unit[sub_txs];
	const int bsh = tx_size_high_unit[sub_txs];
	const int step = bsw * bsh;
	const int row_end =
	AOMMIN(tx_size_high_unit[tx_size], max_blocks_high - blk_row);
	const int col_end =
	AOMMIN(tx_size_wide_unit[tx_size], max_blocks_wide - blk_col);

	assert(bsw > 0 && bsh > 0);

	for (int row = 0; row < row_end; row += bsh) {
	const int offsetr = blk_row + row;
	for (int col = 0; col < col_end; col += bsw) {
	const int offsetc = blk_col + col;

	tokenize_vartx(td, sub_txs, plane_bsize, offsetr, offsetc, block, plane,
	arg);
	block += step;
	}
	}
	}
	}

	void av1_tokenize_sb_vartx(const AV1_COMP cpi, ThreadData td,
	RUN_TYPE dry_run, BLOCK_SIZE bsize, int *rate,
	uint8_t allow_update_cdf) {
	assert(bsize < BLOCK_SIZES_ALL);
	const AV1_COMMON *const cm = &cpi->common;
	MACROBLOCK *const x = &td->mb;
	MACROBLOCKD *const xd = &x->e_mbd;
	const int mi_row = xd->mi_row;
	const int mi_col = xd->mi_col;
	if (mi_row >= cm->mi_params.mi_rows \|\| mi_col >= cm->mi_params.mi_cols)
	return;

	const int num_planes = av1_num_planes(cm);
	MB_MODE_INFO *const mbmi = xd->mi[0];
	struct tokenize_b_args arg = { cpi, td, 0, allow_update_cdf, dry_run };

	if (mbmi->skip_txfm) {
	av1_reset_entropy_context(xd, bsize, num_planes);
	return;
	}

	for (int plane = 0; plane < num_planes; ++plane) {
	if (plane && !xd->is_chroma_ref) break;
	const struct macroblockd_plane *const pd = &xd->plane[plane];
	const int ss_x = pd->subsampling_x;
	const int ss_y = pd->subsampling_y;
	const BLOCK_SIZE plane_bsize = get_plane_block_size(bsize, ss_x, ss_y);
	assert(plane_bsize < BLOCK_SIZES_ALL);
	const int mi_width = mi_size_wide[plane_bsize];
	const int mi_height = mi_size_high[plane_bsize];
	const TX_SIZE max_tx_size = get_vartx_max_txsize(xd, plane_bsize, plane);
	const BLOCK_SIZE txb_size = txsize_to_bsize[max_tx_size];
	const int bw = mi_size_wide[txb_size];
	const int bh = mi_size_high[txb_size];
	int block = 0;
	const int step =
	tx_size_wide_unit[max_tx_size] * tx_size_high_unit[max_tx_size];

	const BLOCK_SIZE max_unit_bsize =
	get_plane_block_size(BLOCK_64X64, ss_x, ss_y);
	int mu_blocks_wide = mi_size_wide[max_unit_bsize];
	int mu_blocks_high = mi_size_high[max_unit_bsize];

	mu_blocks_wide = AOMMIN(mi_width, mu_blocks_wide);
	mu_blocks_high = AOMMIN(mi_height, mu_blocks_high);

	for (int idy = 0; idy < mi_height; idy += mu_blocks_high) {
	for (int idx = 0; idx < mi_width; idx += mu_blocks_wide) {
	const int unit_height = AOMMIN(mu_blocks_high + idy, mi_height);
	const int unit_width = AOMMIN(mu_blocks_wide + idx, mi_width);
	for (int blk_row = idy; blk_row < unit_height; blk_row += bh) {
	for (int blk_col = idx; blk_col < unit_width; blk_col += bw) {
	tokenize_vartx(td, max_tx_size, plane_bsize, blk_row, blk_col,
	block, plane, &arg);
	block += step;
	}
	}
	}
	}
	}
	if (rate) *rate += arg.this_rate;
	}