lavc/aacsbr: fix make checkheaders warning
[ffmpeg.git] / 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 | AV_CODEC_CAP_INTRA_ONLY,
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 };