blob: 580da160010c274213214b64698d9a2e95a05226 [file] [log] [blame]
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include "av1/common/reconinter.h"
#include "av1/common/reconintra.h"
#include "av1/encoder/interintra_ml_data_collect.h"
#include "av1/encoder/reconinter_enc.h"
#define BORDER_SIZE 4 // Only generate data with a border of 4 pixels.
// Static references to the currently captured Y/U/V planes. Held in
// memory until we can determine if this block is a skip-block (if so,
// the data is not written out).
static IIMLPlaneInfo *INFO_Y = NULL;
static IIMLPlaneInfo *INFO_U = NULL;
static IIMLPlaneInfo *INFO_V = NULL;
// Static reference to the file where the data is being stored.
static FILE *FP = NULL;
// Clean-up function called on program exit. Closes file descriptor,
// flushing writes.
static void cleanup_fp() {
if (FP != NULL) {
int r = fclose(FP);
assert(r == 0);
(void)r;
}
}
static uint8_t read_interintra_bias() {
const char *bias = getenv("INTERINTRA_BIAS");
if (bias == NULL) {
return 100;
}
int i = atoi(bias);
if (i < 0 || i > 100) {
fprintf(stderr, "INTERINTRA_BIAS must be between 0 and 100.\n");
exit(1);
}
return (uint8_t)i;
}
static uint8_t INTERINTRA_BIAS = 255;
uint8_t av1_interintra_bias() {
if (INTERINTRA_BIAS == 255) {
INTERINTRA_BIAS = read_interintra_bias();
}
return INTERINTRA_BIAS;
}
// Function called on every invocation of data collection. Initializes
// the file output and registers the cleanup function (if not already done).
// The output file name is "interintra_ml_data_collect.bin" by default.
// The file name can be changed by setting INTERINTRA_ML_DATA_COLLECT
// environment variable.
static void init_first_run() {
if (FP != NULL) {
return;
}
const char *filename = getenv("INTERINTRA_ML_DATA_COLLECT");
if (filename == NULL) {
filename = "interintra_ml_data_collect.bin";
}
FP = fopen(filename, "wb");
assert(FP != NULL);
int r = atexit(cleanup_fp);
assert(r == 0);
(void)r;
}
// Copy the value from the source into the destination, at the given offsets.
// If both source and destination are high-bitdepth, will copy. Otherwise,
// if source is high-bitdepth but destination is lower, will only copy the
// first 8 bits.
static void copy_value(uint8_t *dst, int dst_offset, const uint8_t *src,
int src_offset, int bitdepth, bool is_hbd) {
if (bitdepth > 8) {
assert(is_hbd); // If the IIMLPlaneInfo is high-bitdepth, then so must the
// AV1 data structures.
uint16_t *dst16 = (uint16_t *)dst;
uint16_t *src16 = CONVERT_TO_SHORTPTR(src);
dst16[dst_offset] = src16[src_offset];
} else if (bitdepth == 8 && !is_hbd) {
// A pure copy.
dst[dst_offset] = src[src_offset];
} else {
// There is no case when bitdepth > 8 and !is_hbd.
assert(bitdepth == 8 && is_hbd);
uint16_t *src16 = CONVERT_TO_SHORTPTR(src);
assert(src16[src_offset] <= 255);
dst[dst_offset] = (uint8_t)src16[src_offset];
}
}
static void copy_source_image(IIMLPlaneInfo *info, MACROBLOCK *const x) {
// Overallocate, memory is not an issue.
info->source_image = malloc(info->width * info->height * sizeof(uint16_t));
assert(info->source_image != NULL);
const struct buf_2d *ref = &x->plane[info->plane].src;
const int stride = ref->stride;
const uint8_t *buf = ref->buf0 + info->y * stride + info->x;
const bool is_hbd = is_cur_buf_hbd(&x->e_mbd);
for (int j = 0; j < info->height; ++j) {
for (int i = 0; i < info->width; ++i) {
copy_value(info->source_image, j * info->width + i, buf, j * stride + i,
info->bitdepth, is_hbd);
}
}
}
static void copy_intrapred_lshape(IIMLPlaneInfo *info, MACROBLOCK *const x) {
info->intrapred_lshape =
malloc(sizeof(uint16_t) * ((info->width + info->border) * info->border +
info->border * info->height));
assert(info->intrapred_lshape != NULL);
MACROBLOCKD *const xd = &x->e_mbd;
const struct macroblockd_plane *const pd = &xd->plane[info->plane];
const struct buf_2d *ref = &pd->dst;
const int stride = ref->stride;
uint8_t *buf = ref->buf0 + info->y * stride + info->x;
av1_interintra_ml_data_collect_copy_intrapred_lshape(
info->intrapred_lshape, buf, stride, info->width, info->height,
info->border, info->bitdepth, is_cur_buf_hbd(xd));
}
void av1_interintra_ml_data_collect_copy_intrapred_lshape(
uint8_t *dst, const uint8_t *src, int src_stride, int width, int height,
int border, int bitdepth, bool is_src_hbd) {
// Point back at the start of the L-region.
src -= (border * src_stride + border);
int info_i = 0;
// Copy over the top part.
for (int j = 0; j < border; ++j) {
for (int i = 0; i < border + width; ++i) {
copy_value(dst, info_i++, src, j * src_stride + i, bitdepth, is_src_hbd);
}
}
// Copy over the side pixels.
for (int j = border; j < border + height; ++j) {
for (int i = 0; i < border; ++i) {
copy_value(dst, info_i++, src, j * src_stride + i, bitdepth, is_src_hbd);
}
}
}
static void copy_interpred(IIMLPlaneInfo *info, const AV1_COMP *const cpi,
MACROBLOCK *const x) {
info->interpred = malloc(sizeof(uint16_t) * (info->width + info->border) *
(info->height + info->border));
const AV1_COMMON *const cm = &cpi->common;
MACROBLOCKD *const xd = &x->e_mbd;
uint8_t *orig_buf = xd->plane[info->plane].dst.buf;
int orig_stride = xd->plane[info->plane].dst.stride;
uint8_t *interpred;
int interpred_stride;
const int border =
av1_calc_border(xd, AOM_PLANE_Y, false /* build for obmc */);
assert(border >= BORDER_SIZE);
av1_alloc_buf_with_border(&interpred, &interpred_stride, border,
is_cur_buf_hbd(xd));
xd->plane[info->plane].dst.buf = interpred;
xd->plane[info->plane].dst.stride = interpred_stride;
av1_enc_build_border_only_inter_predictor(cm, xd, xd->mi_row, xd->mi_col,
info->plane);
int info_i = 0;
for (int j = -1 * info->border; j < info->height; ++j) {
for (int i = -1 * info->border; i < info->width; ++i) {
copy_value(info->interpred, info_i++, interpred, j * interpred_stride + i,
info->bitdepth, is_cur_buf_hbd(xd));
}
}
xd->plane[info->plane].dst.buf = orig_buf;
xd->plane[info->plane].dst.stride = orig_stride;
av1_free_buf_with_border(interpred, interpred_stride, border,
is_cur_buf_hbd(xd));
}
static void copy_predictor(IIMLPlaneInfo *info, const AV1_COMP *const cpi,
MACROBLOCK *const x, BLOCK_SIZE bsize) {
info->predictor = malloc(sizeof(uint16_t) * info->width * info->height);
const AV1_COMMON *const cm = &cpi->common;
MACROBLOCKD *const xd = &x->e_mbd;
// Assumes that xd->plane[info->plane].dst.buf can be overwritten in
// the inter-predictor part.
av1_enc_build_inter_predictor(cm, xd, xd->mi_row, xd->mi_col, NULL, bsize,
info->plane, info->plane);
uint8_t *predictor = xd->plane[info->plane].dst.buf;
int stride = xd->plane[info->plane].dst.stride;
int info_i = 0;
for (int j = 0; j < info->height; ++j) {
for (int i = 0; i < info->width; ++i) {
copy_value(info->predictor, info_i++, predictor, j * stride + i,
info->bitdepth, is_cur_buf_hbd(xd));
}
}
}
// 1 == inter, 2 == inter-intra (smooth), 3 == inter-intra (wedge)
static int get_prediction_type(MB_MODE_INFO *mbmi, BLOCK_SIZE bsize) {
if (!is_interintra_pred(mbmi)) {
return 1;
}
const bool wedge_case =
mbmi->use_wedge_interintra && is_interintra_wedge_used(bsize);
return wedge_case ? 3 : 2;
}
// Initializes the IIMLPlaneInfo structure.
static IIMLPlaneInfo *init_plane_info(const AV1_COMP *const cpi,
MACROBLOCK *const x, BLOCK_SIZE bsize,
int plane) {
MACROBLOCKD *const xd = &x->e_mbd;
MB_MODE_INFO *mbmi = xd->mi[0];
const AV1_COMMON *const cm = &cpi->common;
const struct macroblockd_plane *const pd = &xd->plane[plane];
const BLOCK_SIZE plane_bsize =
get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);
assert(is_inter_block(mbmi));
assert(!has_second_ref(mbmi));
assert(!is_intrabc_block(mbmi));
// Data is packed differently for smaller block sizes. Not supported.
assert(block_size_wide[plane_bsize] >= 4);
assert(block_size_high[plane_bsize] >= 4);
IIMLPlaneInfo *info = malloc(sizeof(IIMLPlaneInfo));
assert(info != NULL);
info->width = block_size_wide[plane_bsize];
info->height = block_size_high[plane_bsize];
info->plane = plane;
info->bitdepth = (uint8_t)xd->bd;
info->border = BORDER_SIZE;
info->x = (xd->mi_col * MI_SIZE) >> pd->subsampling_x;
info->y = (xd->mi_row * MI_SIZE) >> pd->subsampling_y;
info->frame_order_hint = (int)cm->current_frame.order_hint;
RefCntBuffer *refbuf = get_ref_frame_buf(cm, xd->mi[0]->ref_frame[0]);
assert(refbuf != NULL);
info->ref_q = refbuf->base_qindex;
info->base_q = cm->base_qindex;
info->lambda = av1_compute_rd_mult_based_on_qindex(cpi, cm->base_qindex);
assert(info->lambda > 0);
assert(info->lambda < 10 * 1000 * 1000);
copy_source_image(info, x);
copy_intrapred_lshape(info, x);
copy_interpred(info, cpi, x);
info->prediction_type = get_prediction_type(mbmi, bsize);
copy_predictor(info, cpi, x, bsize);
info->interintra_bias = av1_interintra_bias();
return info;
}
void write_uint8_and_advance(uint8_t **buf, uint8_t b) {
**buf = b;
(*buf)++;
}
void write_int32_and_advance(uint8_t **buf, int32_t value) {
assert(value >= 0);
for (size_t i = 0; i < sizeof(value); ++i) {
// Little-endian order.
uint8_t byte = 0xff & value;
write_uint8_and_advance(buf, byte);
value >>= 8;
}
}
void write_buffer_and_advance(uint8_t **buf, uint8_t *src, size_t size,
int bitdepth) {
if (bitdepth == 8) {
for (size_t i = 0; i < size; ++i) {
write_uint8_and_advance(buf, src[i]);
}
return;
}
uint16_t *src16 = (uint16_t *)src;
for (size_t i = 0; i < size; ++i) {
// Little-endian order.
uint8_t b1 = (uint8_t)(0xff & src16[i]);
uint8_t b2 = (uint8_t)(0xff & (src16[i] >> 8));
write_uint8_and_advance(buf, b1);
write_uint8_and_advance(buf, b2);
}
}
void av1_interintra_ml_data_collect_serialize(IIMLPlaneInfo *info,
uint8_t **buf, size_t *buf_size) {
const int bytes_per_pixel = info->bitdepth == 8 ? 1 : 2;
const size_t source_size = info->width * info->height;
const size_t intrapred_lshape_size =
((info->border + info->width) * info->border +
info->border * info->height);
const size_t interpred_size =
((info->border + info->width) * (info->border + info->height));
*buf_size =
1 /* width */ + 1 /* height */ + 1 /* plane */ + 1 /* bitdepth */ +
1 /* border */ + sizeof(int32_t) /* x */ + sizeof(int32_t) /* y */ +
sizeof(int32_t) /* frame order hint */ + sizeof(int32_t) /* lambda */ +
1 /* base_q */ + bytes_per_pixel * source_size /* source image */ +
bytes_per_pixel * intrapred_lshape_size /* reconstruction border */ +
1 /* prediction_type */ + bytes_per_pixel * source_size /* predictor */ +
1 /* ref_q */ +
bytes_per_pixel * interpred_size /* inter-predictor + border */ +
1 /* interintra_bias */;
*buf = malloc(*buf_size);
assert(*buf != NULL);
uint8_t *start = *buf;
write_uint8_and_advance(buf, info->width);
write_uint8_and_advance(buf, info->height);
write_uint8_and_advance(buf, info->plane);
write_uint8_and_advance(buf, info->bitdepth);
write_uint8_and_advance(buf, info->border);
write_int32_and_advance(buf, info->x);
write_int32_and_advance(buf, info->y);
write_int32_and_advance(buf, info->frame_order_hint);
write_int32_and_advance(buf, info->lambda);
write_uint8_and_advance(buf, info->base_q);
write_buffer_and_advance(buf, info->source_image, source_size,
info->bitdepth);
write_buffer_and_advance(buf, info->intrapred_lshape, intrapred_lshape_size,
info->bitdepth);
write_uint8_and_advance(buf, info->prediction_type);
write_buffer_and_advance(buf, info->predictor, source_size, info->bitdepth);
write_uint8_and_advance(buf, info->ref_q);
write_buffer_and_advance(buf, info->interpred, interpred_size,
info->bitdepth);
write_uint8_and_advance(buf, info->interintra_bias);
assert(start + *buf_size == *buf);
*buf = start;
}
void av1_interintra_ml_data_collect(const AV1_COMP *const cpi,
MACROBLOCK *const x, BLOCK_SIZE bsize) {
init_first_run();
assert(INFO_Y == NULL);
assert(INFO_U == NULL);
assert(INFO_V == NULL);
INFO_Y = init_plane_info(cpi, x, bsize, AOM_PLANE_Y);
assert(INFO_Y != NULL);
// Sometimes the chroma is derived from the luma. Check if chroma is
// available.
MACROBLOCKD *const xd = &x->e_mbd;
if (xd->mi[0]->chroma_ref_info.is_chroma_ref) {
INFO_U = init_plane_info(cpi, x, bsize, AOM_PLANE_U);
INFO_V = init_plane_info(cpi, x, bsize, AOM_PLANE_V);
assert(INFO_U != NULL);
assert(INFO_V != NULL);
}
}
// Write out the three planes to disk.
void av1_interintra_ml_data_collect_finalize() {
uint8_t *buf;
size_t buf_size;
// If luma is not present, then neither is chroma.
if (INFO_Y == NULL) {
assert(INFO_U == NULL);
assert(INFO_V == NULL);
return;
}
av1_interintra_ml_data_collect_serialize(INFO_Y, &buf, &buf_size);
size_t written = fwrite(buf, sizeof(*buf), buf_size, FP);
assert(written == buf_size);
free(buf);
if (INFO_U != NULL) {
av1_interintra_ml_data_collect_serialize(INFO_U, &buf, &buf_size);
written = fwrite(buf, sizeof(*buf), buf_size, FP);
assert(written == buf_size);
free(buf);
}
if (INFO_V != NULL) {
av1_interintra_ml_data_collect_serialize(INFO_V, &buf, &buf_size);
written = fwrite(buf, sizeof(*buf), buf_size, FP);
assert(written == buf_size);
free(buf);
}
av1_interintra_ml_data_collect_abandon();
assert(INFO_Y == NULL);
assert(INFO_U == NULL);
assert(INFO_V == NULL);
}
static void destroy_info(IIMLPlaneInfo *info) {
if (info != NULL) {
free(info->source_image);
free(info->intrapred_lshape);
free(info->interpred);
free(info);
}
}
void av1_interintra_ml_data_collect_abandon() {
if (INFO_Y == NULL) {
assert(INFO_U == NULL);
assert(INFO_V == NULL);
return;
}
destroy_info(INFO_Y);
destroy_info(INFO_U);
destroy_info(INFO_V);
INFO_Y = NULL;
INFO_U = NULL;
INFO_V = NULL;
}
static int calc_intra_border(MACROBLOCK *const x, BLOCK_SIZE bsize) {
MACROBLOCKD *const xd = &x->e_mbd;
int border = BORDER_SIZE;
// For each plane, check how much is available on the top and left,
// shrinking the border if needed.
for (int plane = AOM_PLANE_Y; plane <= AOM_PLANE_V; ++plane) {
border = AOMMIN(border, av1_intra_top_available(xd, plane));
border = AOMMIN(border, av1_intra_left_available(xd, plane));
// If either the bottom or the right side of the block is partially
// unavailable, set the entire border to 0 -- we cannot extract the
// full border.
const struct macroblockd_plane *const pd = &xd->plane[plane];
const BLOCK_SIZE plane_bsize =
get_plane_block_size(bsize, pd->subsampling_x, pd->subsampling_y);
const TX_SIZE tx_size = max_txsize_rect_lookup[plane_bsize];
if (av1_intra_bottom_unavailable(xd, plane, tx_size) ||
av1_intra_right_unavailable(xd, plane, tx_size)) {
return 0;
}
}
return border;
}
bool av1_interintra_ml_data_collect_valid(MACROBLOCK *const x,
BLOCK_SIZE bsize) {
MACROBLOCKD *const xd = &x->e_mbd;
const int inter_border =
av1_calc_border(xd, AOM_PLANE_Y, false /* build for obmc */);
const int intra_border = calc_intra_border(x, bsize);
// Only process 8x8 or larger blocks -- if a dimension is 4, then the
// chroma dimension can be 2, which has a different packing structure.
return inter_border >= BORDER_SIZE && intra_border >= BORDER_SIZE &&
!x->skip && block_size_wide[bsize] >= 8 && block_size_high[bsize] >= 8;
}