git.sesse.net Git - ffmpeg/blob - libavcodec/magicyuvenc.c

   1 /*
   2  * MagicYUV encoder
   3  * Copyright (c) 2017 Paul B Mahol
   4  *
   5  * This file is part of FFmpeg.
   6  *
   7  * FFmpeg is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU Lesser General Public
   9  * License as published by the Free Software Foundation; either
  10  * version 2.1 of the License, or (at your option) any later version.
  11  *
  12  * FFmpeg is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15  * Lesser General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU Lesser General Public
  18  * License along with FFmpeg; if not, write to the Free Software
  19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  20  */
  21
  22 #include <stdlib.h>
  23 #include <string.h>
  24
  25 #include "libavutil/opt.h"
  26 #include "libavutil/pixdesc.h"
  27 #include "libavutil/qsort.h"
  28
  29 #include "avcodec.h"
  30 #include "bytestream.h"
  31 #include "put_bits.h"
  32 #include "internal.h"
  33 #include "thread.h"
  34 #include "lossless_videoencdsp.h"
  35
  36 typedef enum Prediction {
  37     LEFT = 1,
  38     GRADIENT,
  39     MEDIAN,
  40 } Prediction;
  41
  42 typedef struct HuffEntry {
  43     uint8_t  sym;
  44     uint8_t  len;
  45     uint32_t code;
  46 } HuffEntry;
  47
  48 typedef struct PTable {
  49     int     value;  ///< input value
  50     int64_t prob;   ///< number of occurences of this value in input
  51 } PTable;
  52
  53 typedef struct MagicYUVContext {
  54     const AVClass       *class;
  55     int                  frame_pred;
  56     PutBitContext        pb;
  57     int                  planes;
  58     uint8_t              format;
  59     AVFrame             *p;
  60     int                  slice_height;
  61     int                  nb_slices;
  62     int                  correlate;
  63     int                  hshift[4];
  64     int                  vshift[4];
  65     uint8_t             *slices[4];
  66     unsigned             slice_pos[4];
  67     unsigned             tables_size;
  68     HuffEntry            he[4][256];
  69     LLVidEncDSPContext   llvidencdsp;
  70     void (*predict)(struct MagicYUVContext *s, uint8_t *src, uint8_t *dst,
  71                     ptrdiff_t stride, int width, int height);
  72 } MagicYUVContext;
  73
  74 static void left_predict(MagicYUVContext *s,
  75                          uint8_t *src, uint8_t *dst, ptrdiff_t stride,
  76                          int width, int height)
  77 {
  78     uint8_t prev = 0;
  79     int i, j;
  80
  81     for (i = 0; i < width; i++) {
  82         dst[i] = src[i] - prev;
  83         prev   = src[i];
  84     }
  85     dst += width;
  86     src += stride;
  87     for (j = 1; j < height; j++) {
  88         prev = src[-stride];
  89         for (i = 0; i < width; i++) {
  90             dst[i] = src[i] - prev;
  91             prev   = src[i];
  92         }
  93         dst += width;
  94         src += stride;
  95     }
  96 }
  97
  98 static void gradient_predict(MagicYUVContext *s,
  99                              uint8_t *src, uint8_t *dst, ptrdiff_t stride,
 100                              int width, int height)
 101 {
 102     int left = 0, top, lefttop;
 103     int i, j;
 104
 105     for (i = 0; i < width; i++) {
 106         dst[i] = src[i] - left;
 107         left   = src[i];
 108     }
 109     dst += width;
 110     src += stride;
 111     for (j = 1; j < height; j++) {
 112         top = src[-stride];
 113         left = src[0] - top;
 114         dst[0] = left;
 115         for (i = 1; i < width; i++) {
 116             top = src[i - stride];
 117             lefttop = src[i - (stride + 1)];
 118             left = src[i-1];
 119             dst[i] = (src[i] - top) - left + lefttop;
 120         }
 121         dst += width;
 122         src += stride;
 123     }
 124 }
 125
 126 static void median_predict(MagicYUVContext *s,
 127                            uint8_t *src, uint8_t *dst, ptrdiff_t stride,
 128                            int width, int height)
 129 {
 130     int left = 0, lefttop;
 131     int i, j;
 132
 133     for (i = 0; i < width; i++) {
 134         dst[i] = src[i] - left;
 135         left   = src[i];
 136     }
 137     dst += width;
 138     src += stride;
 139     for (j = 1; j < height; j++) {
 140         left = lefttop = src[-stride];
 141         s->llvidencdsp.sub_median_pred(dst, src - stride, src, width, &left, &lefttop);
 142         dst += width;
 143         src += stride;
 144     }
 145 }
 146
 147 static av_cold int magy_encode_init(AVCodecContext *avctx)
 148 {
 149     MagicYUVContext *s = avctx->priv_data;
 150     PutByteContext pb;
 151     int i;
 152
 153     switch (avctx->pix_fmt) {
 154     case AV_PIX_FMT_GBRP:
 155         avctx->codec_tag = MKTAG('M', '8', 'R', 'G');
 156         s->correlate = 1;
 157         s->format = 0x65;
 158         break;
 159     case AV_PIX_FMT_GBRAP:
 160         avctx->codec_tag = MKTAG('M', '8', 'R', 'A');
 161         s->correlate = 1;
 162         s->format = 0x66;
 163         break;
 164     case AV_PIX_FMT_YUV420P:
 165         avctx->codec_tag = MKTAG('M', '8', 'Y', '0');
 166         s->hshift[1] =
 167         s->vshift[1] =
 168         s->hshift[2] =
 169         s->vshift[2] = 1;
 170         s->format = 0x69;
 171         break;
 172     case AV_PIX_FMT_YUV422P:
 173         avctx->codec_tag = MKTAG('M', '8', 'Y', '2');
 174         s->hshift[1] =
 175         s->hshift[2] = 1;
 176         s->format = 0x68;
 177         break;
 178     case AV_PIX_FMT_YUV444P:
 179         avctx->codec_tag = MKTAG('M', '8', 'Y', '4');
 180         s->format = 0x67;
 181         break;
 182     case AV_PIX_FMT_YUVA444P:
 183         avctx->codec_tag = MKTAG('M', '8', 'Y', 'A');
 184         s->format = 0x6a;
 185         break;
 186     case AV_PIX_FMT_GRAY8:
 187         avctx->codec_tag = MKTAG('M', '8', 'G', '0');
 188         s->format = 0x6b;
 189         break;
 190     default:
 191         av_log(avctx, AV_LOG_ERROR, "Unsupported pixel format: %d\n",
 192                avctx->pix_fmt);
 193         return AVERROR_INVALIDDATA;
 194     }
 195
 196     ff_llvidencdsp_init(&s->llvidencdsp);
 197
 198     s->planes = av_pix_fmt_count_planes(avctx->pix_fmt);
 199
 200     s->nb_slices = 1;
 201
 202     for (i = 0; i < s->planes; i++) {
 203         s->slices[i] = av_malloc(avctx->width * (avctx->height + 2) +
 204                                  AV_INPUT_BUFFER_PADDING_SIZE);
 205         if (!s->slices[i]) {
 206             av_log(avctx, AV_LOG_ERROR, "Cannot allocate temporary buffer.\n");
 207             return AVERROR(ENOMEM);
 208         }
 209     }
 210
 211     switch (s->frame_pred) {
 212     case LEFT:     s->predict = left_predict;     break;
 213     case GRADIENT: s->predict = gradient_predict; break;
 214     case MEDIAN:   s->predict = median_predict;   break;
 215     }
 216
 217     avctx->extradata_size = 32;
 218
 219     avctx->extradata = av_mallocz(avctx->extradata_size +
 220                                   AV_INPUT_BUFFER_PADDING_SIZE);
 221
 222     if (!avctx->extradata) {
 223         av_log(avctx, AV_LOG_ERROR, "Could not allocate extradata.\n");
 224         return AVERROR(ENOMEM);
 225     }
 226
 227     bytestream2_init_writer(&pb, avctx->extradata, avctx->extradata_size);
 228     bytestream2_put_le32(&pb, MKTAG('M', 'A', 'G', 'Y'));
 229     bytestream2_put_le32(&pb, 32);
 230     bytestream2_put_byte(&pb, 7);
 231     bytestream2_put_byte(&pb, s->format);
 232     bytestream2_put_byte(&pb, 12);
 233     bytestream2_put_byte(&pb, 0);
 234
 235     bytestream2_put_byte(&pb, 0);
 236     bytestream2_put_byte(&pb, 0);
 237     bytestream2_put_byte(&pb, 32);
 238     bytestream2_put_byte(&pb, 0);
 239
 240     bytestream2_put_le32(&pb, avctx->width);
 241     bytestream2_put_le32(&pb, avctx->height);
 242     bytestream2_put_le32(&pb, avctx->width);
 243     bytestream2_put_le32(&pb, avctx->height);
 244
 245     return 0;
 246 }
 247
 248 static int magy_huff_cmp_len(const void *a, const void *b)
 249 {
 250     const HuffEntry *aa = a, *bb = b;
 251     return (aa->len - bb->len) * 256 + aa->sym - bb->sym;
 252 }
 253
 254 static int huff_cmp_sym(const void *a, const void *b)
 255 {
 256     const HuffEntry *aa = a, *bb = b;
 257     return bb->sym - aa->sym;
 258 }
 259
 260 static void calculate_codes(HuffEntry *he)
 261 {
 262     uint32_t code;
 263     int i;
 264
 265     AV_QSORT(he, 256, HuffEntry, magy_huff_cmp_len);
 266
 267     code = 1;
 268     for (i = 255; i >= 0; i--) {
 269         he[i].code  = code >> (32 - he[i].len);
 270         code       += 0x80000000u >> (he[i].len - 1);
 271     }
 272
 273     AV_QSORT(he, 256, HuffEntry, huff_cmp_sym);
 274 }
 275
 276 static void count_usage(uint8_t *src, int width,
 277                         int height, PTable *counts)
 278 {
 279     int i, j;
 280
 281     for (j = 0; j < height; j++) {
 282         for (i = 0; i < width; i++) {
 283             counts[src[i]].prob++;
 284         }
 285         src += width;
 286     }
 287 }
 288
 289 typedef struct PackageMergerList {
 290     int nitems;             ///< number of items in the list and probability      ex. 4
 291     int item_idx[515];      ///< index range for each item in items                   0, 2, 5, 9, 13
 292     int probability[514];   ///< probability of each item                             3, 8, 18, 46
 293     int items[257 * 16];    ///< chain of all individual values that make up items    A, B, A, B, C, A, B, C, D, C, D, D, E
 294 } PackageMergerList;
 295
 296 static int compare_by_prob(const void *a, const void *b)
 297 {
 298     PTable a_val = *(PTable *)a;
 299     PTable b_val = *(PTable *)b;
 300     return a_val.prob - b_val.prob;
 301 }
 302
 303 static void magy_huffman_compute_bits(PTable *prob_table, HuffEntry *distincts,
 304                                       int size, int max_length)
 305 {
 306     PackageMergerList list_a, list_b, *to = &list_a, *from = &list_b, *temp;
 307     int times, i, j, k;
 308     int nbits[257] = {0};
 309     int min;
 310
 311     av_assert0(max_length > 0);
 312
 313     to->nitems = 0;
 314     from->nitems = 0;
 315     to->item_idx[0] = 0;
 316     from->item_idx[0] = 0;
 317     AV_QSORT(prob_table, size, PTable, compare_by_prob);
 318
 319     for (times = 0; times <= max_length; times++) {
 320         to->nitems = 0;
 321         to->item_idx[0] = 0;
 322
 323         j = 0;
 324         k = 0;
 325
 326         if (times < max_length) {
 327             i = 0;
 328         }
 329         while (i < size || j + 1 < from->nitems) {
 330             to->nitems++;
 331             to->item_idx[to->nitems] = to->item_idx[to->nitems - 1];
 332             if (i < size &&
 333                 (j + 1 >= from->nitems ||
 334                  prob_table[i].prob <
 335                      from->probability[j] + from->probability[j + 1])) {
 336                 to->items[to->item_idx[to->nitems]++] = prob_table[i].value;
 337                 to->probability[to->nitems - 1] = prob_table[i].prob;
 338                 i++;
 339             } else {
 340                 for (k = from->item_idx[j]; k < from->item_idx[j + 2]; k++) {
 341                     to->items[to->item_idx[to->nitems]++] = from->items[k];
 342                 }
 343                 to->probability[to->nitems - 1] =
 344                     from->probability[j] + from->probability[j + 1];
 345                 j += 2;
 346             }
 347         }
 348         temp = to;
 349         to = from;
 350         from = temp;
 351     }
 352
 353     min = (size - 1 < from->nitems) ? size - 1 : from->nitems;
 354     for (i = 0; i < from->item_idx[min]; i++) {
 355         nbits[from->items[i]]++;
 356     }
 357
 358     for (i = 0; i < size; i++) {
 359         distincts[i].sym = i;
 360         distincts[i].len = nbits[i];
 361     }
 362 }
 363
 364 static int encode_table(AVCodecContext *avctx, uint8_t *dst,
 365                         int width, int height,
 366                         PutBitContext *pb, HuffEntry *he)
 367 {
 368     PTable counts[256] = { {0} };
 369     int i;
 370
 371     count_usage(dst, width, height, counts);
 372
 373     for (i = 0; i < 256; i++) {
 374         counts[i].prob++;
 375         counts[i].value = 255 - i;
 376     }
 377
 378     magy_huffman_compute_bits(counts, he, 256, 12);
 379
 380     calculate_codes(he);
 381
 382     for (i = 0; i < 256; i++) {
 383         put_bits(pb, 1, 0);
 384         put_bits(pb, 7, he[i].len);
 385     }
 386
 387     return 0;
 388 }
 389
 390 static int encode_slice(uint8_t *src, uint8_t *dst, int dst_size,
 391                         int width, int height, HuffEntry *he, int prediction)
 392 {
 393     PutBitContext pb;
 394     int i, j;
 395     int count;
 396
 397     init_put_bits(&pb, dst, dst_size);
 398
 399     put_bits(&pb, 8, 0);
 400     put_bits(&pb, 8, prediction);
 401
 402     for (j = 0; j < height; j++) {
 403         for (i = 0; i < width; i++) {
 404             const int idx = src[i];
 405             put_bits(&pb, he[idx].len, he[idx].code);
 406         }
 407
 408         src += width;
 409     }
 410
 411     count = put_bits_count(&pb) & 0x1F;
 412
 413     if (count)
 414         put_bits(&pb, 32 - count, 0);
 415
 416     count = put_bits_count(&pb);
 417
 418     flush_put_bits(&pb);
 419
 420     return count >> 3;
 421 }
 422
 423 static int magy_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
 424                              const AVFrame *frame, int *got_packet)
 425 {
 426     MagicYUVContext *s = avctx->priv_data;
 427     PutByteContext pb;
 428     const int width = avctx->width, height = avctx->height;
 429     int pos, slice, i, j, ret = 0;
 430
 431     ret = ff_alloc_packet2(avctx, pkt, (256 + 4 * s->nb_slices + width * height) *
 432                            s->planes + 256, 0);
 433     if (ret < 0)
 434         return ret;
 435
 436     bytestream2_init_writer(&pb, pkt->data, pkt->size);
 437     bytestream2_put_le32(&pb, MKTAG('M', 'A', 'G', 'Y'));
 438     bytestream2_put_le32(&pb, 32); // header size
 439     bytestream2_put_byte(&pb, 7);  // version
 440     bytestream2_put_byte(&pb, s->format);
 441     bytestream2_put_byte(&pb, 12); // max huffman length
 442     bytestream2_put_byte(&pb, 0);
 443
 444     bytestream2_put_byte(&pb, 0);
 445     bytestream2_put_byte(&pb, 0);
 446     bytestream2_put_byte(&pb, 32); // coder type
 447     bytestream2_put_byte(&pb, 0);
 448
 449     bytestream2_put_le32(&pb, avctx->width);
 450     bytestream2_put_le32(&pb, avctx->height);
 451     bytestream2_put_le32(&pb, avctx->width);
 452     bytestream2_put_le32(&pb, avctx->height);
 453     bytestream2_put_le32(&pb, 0);
 454
 455     for (i = 0; i < s->planes; i++) {
 456         bytestream2_put_le32(&pb, 0);
 457         for (j = 1; j < s->nb_slices; j++) {
 458             bytestream2_put_le32(&pb, 0);
 459         }
 460     }
 461
 462     bytestream2_put_byte(&pb, s->planes);
 463
 464     for (i = 0; i < s->planes; i++) {
 465         for (slice = 0; slice < s->nb_slices; slice++) {
 466             bytestream2_put_byte(&pb, i);
 467         }
 468     }
 469
 470     if (s->correlate) {
 471         uint8_t *r, *g, *b;
 472         AVFrame *p = av_frame_clone(frame);
 473
 474         g = p->data[0];
 475         b = p->data[1];
 476         r = p->data[2];
 477
 478         for (i = 0; i < height; i++) {
 479             s->llvidencdsp.diff_bytes(b, b, g, width);
 480             s->llvidencdsp.diff_bytes(r, r, g, width);
 481             g += p->linesize[0];
 482             b += p->linesize[1];
 483             r += p->linesize[2];
 484         }
 485
 486         FFSWAP(uint8_t*, p->data[0], p->data[1]);
 487         FFSWAP(int, p->linesize[0], p->linesize[1]);
 488
 489         for (i = 0; i < s->planes; i++) {
 490             for (slice = 0; slice < s->nb_slices; slice++) {
 491                 s->predict(s, p->data[i], s->slices[i], p->linesize[i],
 492                                p->width, p->height);
 493             }
 494         }
 495
 496         av_frame_free(&p);
 497     } else {
 498         for (i = 0; i < s->planes; i++) {
 499             for (slice = 0; slice < s->nb_slices; slice++) {
 500                 s->predict(s, frame->data[i], s->slices[i], frame->linesize[i],
 501                            AV_CEIL_RSHIFT(frame->width, s->hshift[i]),
 502                            AV_CEIL_RSHIFT(frame->height, s->vshift[i]));
 503             }
 504         }
 505     }
 506
 507     init_put_bits(&s->pb, pkt->data + bytestream2_tell_p(&pb), bytestream2_get_bytes_left_p(&pb));
 508
 509     for (i = 0; i < s->planes; i++) {
 510         encode_table(avctx, s->slices[i],
 511                      AV_CEIL_RSHIFT(frame->width,  s->hshift[i]),
 512                      AV_CEIL_RSHIFT(frame->height, s->vshift[i]),
 513                      &s->pb, s->he[i]);
 514     }
 515     s->tables_size = (put_bits_count(&s->pb) + 7) >> 3;
 516     bytestream2_skip_p(&pb, s->tables_size);
 517
 518     for (i = 0; i < s->planes; i++) {
 519         unsigned slice_size;
 520
 521         s->slice_pos[i] = bytestream2_tell_p(&pb);
 522         slice_size = encode_slice(s->slices[i], pkt->data + bytestream2_tell_p(&pb),
 523                                   bytestream2_get_bytes_left_p(&pb),
 524                                   AV_CEIL_RSHIFT(frame->width,  s->hshift[i]),
 525                                   AV_CEIL_RSHIFT(frame->height, s->vshift[i]),
 526                                   s->he[i], s->frame_pred);
 527         bytestream2_skip_p(&pb, slice_size);
 528     }
 529
 530     pos = bytestream2_tell_p(&pb);
 531     bytestream2_seek_p(&pb, 32, SEEK_SET);
 532     bytestream2_put_le32(&pb, s->slice_pos[0] - 32);
 533     for (i = 0; i < s->planes; i++) {
 534         bytestream2_put_le32(&pb, s->slice_pos[i] - 32);
 535     }
 536     bytestream2_seek_p(&pb, pos, SEEK_SET);
 537
 538     pkt->size   = bytestream2_tell_p(&pb);
 539     pkt->flags |= AV_PKT_FLAG_KEY;
 540
 541     *got_packet = 1;
 542
 543     return 0;
 544 }
 545
 546 static av_cold int magy_encode_close(AVCodecContext *avctx)
 547 {
 548     MagicYUVContext *s = avctx->priv_data;
 549     int i;
 550
 551     for (i = 0; i < s->planes; i++)
 552         av_freep(&s->slices[i]);
 553
 554     return 0;
 555 }
 556
 557 #define OFFSET(x) offsetof(MagicYUVContext, x)
 558 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
 559 static const AVOption options[] = {
 560     { "pred", "Prediction method", OFFSET(frame_pred), AV_OPT_TYPE_INT, {.i64=LEFT}, LEFT, MEDIAN, VE, "pred" },
 561     { "left",     NULL, 0, AV_OPT_TYPE_CONST, { .i64 = LEFT },     0, 0, VE, "pred" },
 562     { "gradient", NULL, 0, AV_OPT_TYPE_CONST, { .i64 = GRADIENT }, 0, 0, VE, "pred" },
 563     { "median",   NULL, 0, AV_OPT_TYPE_CONST, { .i64 = MEDIAN },   0, 0, VE, "pred" },
 564     { NULL},
 565 };
 566
 567 static const AVClass magicyuv_class = {
 568     .class_name = "magicyuv",
 569     .item_name  = av_default_item_name,
 570     .option     = options,
 571     .version    = LIBAVUTIL_VERSION_INT,
 572 };
 573
 574 AVCodec ff_magicyuv_encoder = {
 575     .name             = "magicyuv",
 576     .long_name        = NULL_IF_CONFIG_SMALL("MagicYUV video"),
 577     .type             = AVMEDIA_TYPE_VIDEO,
 578     .id               = AV_CODEC_ID_MAGICYUV,
 579     .priv_data_size   = sizeof(MagicYUVContext),
 580     .priv_class       = &magicyuv_class,
 581     .init             = magy_encode_init,
 582     .close            = magy_encode_close,
 583     .encode2          = magy_encode_frame,
 584     .capabilities     = AV_CODEC_CAP_FRAME_THREADS,
 585     .pix_fmts         = (const enum AVPixelFormat[]) {
 586                           AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRAP, AV_PIX_FMT_YUV422P,
 587                           AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUVA444P, AV_PIX_FMT_GRAY8,
 588                           AV_PIX_FMT_NONE
 589                       },
 590 };