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
left_predict(MagicYUVContext * s,uint8_t * src,uint8_t * dst,ptrdiff_t stride,int width,int height)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
gradient_predict(MagicYUVContext * s,uint8_t * src,uint8_t * dst,ptrdiff_t stride,int width,int height)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
median_predict(MagicYUVContext * s,uint8_t * src,uint8_t * dst,ptrdiff_t stride,int width,int height)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
magy_encode_init(AVCodecContext * avctx)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
calculate_codes(HuffEntry * he,uint16_t codes_count[33])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
count_usage(uint8_t * src,int width,int height,PTable * counts)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
compare_by_prob(const void * a,const void * b)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
magy_huffman_compute_bits(PTable * prob_table,HuffEntry * distincts,uint16_t codes_counts[33],int size,int max_length)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
encode_table(AVCodecContext * avctx,uint8_t * dst,int width,int height,PutBitContext * pb,HuffEntry * he)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
encode_slice(uint8_t * src,uint8_t * dst,int dst_size,int width,int height,HuffEntry * he,int prediction)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
magy_encode_frame(AVCodecContext * avctx,AVPacket * pkt,const AVFrame * frame,int * got_packet)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
magy_encode_close(AVCodecContext * avctx)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 };
579