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