]> git.sesse.net Git - ffmpeg/blob - libavcodec/exrenc.c
avcodec/exrenc: add half-float support
[ffmpeg] / libavcodec / exrenc.c
1 /*
2  * Copyright (c) 2021 Paul B Mahol
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or
7  * modify it under the terms of the GNU Lesser General Public
8  * License as published by the Free Software Foundation; either
9  * version 2.1 of the License, or (at your option) any later version.
10  *
11  * FFmpeg is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14  * Lesser General Public License for more details.
15  *
16  * You should have received a copy of the GNU Lesser General Public
17  * License along with FFmpeg; if not, write to the Free Software
18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19  */
20
21 /**
22  * @file
23  * OpenEXR encoder
24  */
25
26 #include <float.h>
27 #include <zlib.h>
28
29 #include "libavutil/avassert.h"
30 #include "libavutil/opt.h"
31 #include "libavutil/intreadwrite.h"
32 #include "libavutil/imgutils.h"
33 #include "libavutil/pixdesc.h"
34 #include "avcodec.h"
35 #include "bytestream.h"
36 #include "internal.h"
37
38 enum ExrCompr {
39     EXR_RAW,
40     EXR_RLE,
41     EXR_ZIP1,
42     EXR_ZIP16,
43     EXR_NBCOMPR,
44 };
45
46 enum ExrPixelType {
47     EXR_UINT,
48     EXR_HALF,
49     EXR_FLOAT,
50     EXR_UNKNOWN,
51 };
52
53 static const char abgr_chlist[4] = { 'A', 'B', 'G', 'R' };
54 static const char bgr_chlist[4] = { 'B', 'G', 'R', 'A' };
55 static const uint8_t gbra_order[4] = { 3, 1, 0, 2 };
56 static const uint8_t gbr_order[4] = { 1, 0, 2, 0 };
57
58 typedef struct EXRScanlineData {
59     uint8_t *compressed_data;
60     unsigned int compressed_size;
61
62     uint8_t *uncompressed_data;
63     unsigned int uncompressed_size;
64
65     uint8_t *tmp;
66     unsigned int tmp_size;
67
68     int64_t actual_size;
69 } EXRScanlineData;
70
71 typedef struct EXRContext {
72     const AVClass *class;
73
74     int compression;
75     int pixel_type;
76     int planes;
77     int nb_scanlines;
78     int scanline_height;
79     float gamma;
80     const char *ch_names;
81     const uint8_t *ch_order;
82     PutByteContext pb;
83
84     EXRScanlineData *scanline;
85
86     uint16_t basetable[512];
87     uint8_t shifttable[512];
88 } EXRContext;
89
90 static void half_tables(EXRContext *s)
91 {
92     for (int i = 0; i < 256; i++) {
93         int e = i - 127;
94
95         if (e < -24) { // Very small numbers map to zero
96             s->basetable[i|0x000]  = 0x0000;
97             s->basetable[i|0x100]  = 0x8000;
98             s->shifttable[i|0x000] = 24;
99             s->shifttable[i|0x100] = 24;
100         } else if (e < -14) { // Small numbers map to denorms
101             s->basetable[i|0x000] = (0x0400>>(-e-14));
102             s->basetable[i|0x100] = (0x0400>>(-e-14)) | 0x8000;
103             s->shifttable[i|0x000] = -e-1;
104             s->shifttable[i|0x100] = -e-1;
105         } else if (e <= 15) { // Normal numbers just lose precision
106             s->basetable[i|0x000] = ((e + 15) << 10);
107             s->basetable[i|0x100] = ((e + 15) << 10) | 0x8000;
108             s->shifttable[i|0x000] = 13;
109             s->shifttable[i|0x100] = 13;
110         } else if (e < 128) { // Large numbers map to Infinity
111             s->basetable[i|0x000]  = 0x7C00;
112             s->basetable[i|0x100]  = 0xFC00;
113             s->shifttable[i|0x000] = 24;
114             s->shifttable[i|0x100] = 24;
115         } else{ // Infinity and NaN's stay Infinity and NaN's
116             s->basetable[i|0x000]  = 0x7C00;
117             s->basetable[i|0x100]  = 0xFC00;
118             s->shifttable[i|0x000] = 13;
119             s->shifttable[i|0x100] = 13;
120         }
121     }
122 }
123
124 static uint16_t float2half(EXRContext *s, uint32_t f)
125 {
126     uint16_t h;
127
128     h = s->basetable[(f >> 23) & 0x1ff] + ((f & 0x007fffff) >> s->shifttable[(f >> 23) & 0x1ff]);
129
130     return h;
131 }
132
133 static int encode_init(AVCodecContext *avctx)
134 {
135     EXRContext *s = avctx->priv_data;
136
137     half_tables(s);
138
139     switch (avctx->pix_fmt) {
140     case AV_PIX_FMT_GBRPF32:
141         s->planes = 3;
142         s->ch_names = bgr_chlist;
143         s->ch_order = gbr_order;
144         break;
145     case AV_PIX_FMT_GBRAPF32:
146         s->planes = 4;
147         s->ch_names = abgr_chlist;
148         s->ch_order = gbra_order;
149         break;
150     default:
151         av_assert0(0);
152     }
153
154     switch (s->compression) {
155     case EXR_RAW:
156     case EXR_RLE:
157     case EXR_ZIP1:
158         s->scanline_height = 1;
159         s->nb_scanlines = avctx->height;
160         break;
161     case EXR_ZIP16:
162         s->scanline_height = 16;
163         s->nb_scanlines = (avctx->height + s->scanline_height - 1) / s->scanline_height;
164         break;
165     default:
166         av_assert0(0);
167     }
168
169     s->scanline = av_calloc(s->nb_scanlines, sizeof(*s->scanline));
170     if (!s->scanline)
171         return AVERROR(ENOMEM);
172
173     return 0;
174 }
175
176 static int encode_close(AVCodecContext *avctx)
177 {
178     EXRContext *s = avctx->priv_data;
179
180     for (int y = 0; y < s->nb_scanlines && s->scanline; y++) {
181         EXRScanlineData *scanline = &s->scanline[y];
182
183         av_freep(&scanline->tmp);
184         av_freep(&scanline->compressed_data);
185         av_freep(&scanline->uncompressed_data);
186     }
187
188     av_freep(&s->scanline);
189
190     return 0;
191 }
192
193 static void reorder_pixels(uint8_t *dst, const uint8_t *src, ptrdiff_t size)
194 {
195     const ptrdiff_t half_size = (size + 1) / 2;
196     uint8_t *t1 = dst;
197     uint8_t *t2 = dst + half_size;
198
199     for (ptrdiff_t i = 0; i < half_size; i++) {
200         t1[i] = *(src++);
201         t2[i] = *(src++);
202     }
203 }
204
205 static void predictor(uint8_t *src, ptrdiff_t size)
206 {
207     int p = src[0];
208
209     for (ptrdiff_t i = 1; i < size; i++) {
210         int d = src[i] - p + 384;
211
212         p = src[i];
213         src[i] = d;
214     }
215 }
216
217 static int64_t rle_compress(uint8_t *out, int64_t out_size,
218                             const uint8_t *in, int64_t in_size)
219 {
220     int64_t i = 0, o = 0, run = 1, copy = 0;
221
222     while (i < in_size) {
223         while (i + run < in_size && in[i] == in[i + run] && run < 128)
224             run++;
225
226         if (run >= 3) {
227             if (o + 2 >= out_size)
228                 return -1;
229             out[o++] = run - 1;
230             out[o++] = in[i];
231             i += run;
232         } else {
233             if (i + run < in_size)
234                 copy += run;
235             while (i + copy < in_size && copy < 127 && in[i + copy] != in[i + copy - 1])
236                 copy++;
237
238             if (o + 1 + copy >= out_size)
239                 return -1;
240             out[o++] = -copy;
241
242             for (int x = 0; x < copy; x++)
243                 out[o + x] = in[i + x];
244
245             o += copy;
246             i += copy;
247             copy = 0;
248         }
249
250         run = 1;
251     }
252
253     return o;
254 }
255
256 static int encode_scanline_rle(EXRContext *s, const AVFrame *frame)
257 {
258     const int64_t element_size = s->pixel_type == EXR_HALF ? 2LL : 4LL;
259
260     for (int y = 0; y < frame->height; y++) {
261         EXRScanlineData *scanline = &s->scanline[y];
262         int64_t tmp_size = element_size * s->planes * frame->width;
263         int64_t max_compressed_size = tmp_size * 3 / 2;
264
265         av_fast_padded_malloc(&scanline->uncompressed_data, &scanline->uncompressed_size, tmp_size);
266         if (!scanline->uncompressed_data)
267             return AVERROR(ENOMEM);
268
269         av_fast_padded_malloc(&scanline->tmp, &scanline->tmp_size, tmp_size);
270         if (!scanline->tmp)
271             return AVERROR(ENOMEM);
272
273         av_fast_padded_malloc(&scanline->compressed_data, &scanline->compressed_size, max_compressed_size);
274         if (!scanline->compressed_data)
275             return AVERROR(ENOMEM);
276
277         switch (s->pixel_type) {
278         case EXR_FLOAT:
279             for (int p = 0; p < s->planes; p++) {
280                 int ch = s->ch_order[p];
281
282                 memcpy(scanline->uncompressed_data + frame->width * 4 * p,
283                        frame->data[ch] + y * frame->linesize[ch], frame->width * 4);
284             }
285             break;
286         case EXR_HALF:
287             for (int p = 0; p < s->planes; p++) {
288                 int ch = s->ch_order[p];
289                 uint16_t *dst = (uint16_t *)(scanline->uncompressed_data + frame->width * 2 * p);
290                 uint32_t *src = (uint32_t *)(frame->data[ch] + y * frame->linesize[ch]);
291
292                 for (int x = 0; x < frame->width; x++)
293                     dst[x] = float2half(s, src[x]);
294             }
295             break;
296         }
297
298         reorder_pixels(scanline->tmp, scanline->uncompressed_data, tmp_size);
299         predictor(scanline->tmp, tmp_size);
300         scanline->actual_size = rle_compress(scanline->compressed_data,
301                                              max_compressed_size,
302                                              scanline->tmp, tmp_size);
303
304         if (scanline->actual_size <= 0 || scanline->actual_size >= tmp_size) {
305             FFSWAP(uint8_t *, scanline->uncompressed_data, scanline->compressed_data);
306             FFSWAP(int, scanline->uncompressed_size, scanline->compressed_size);
307             scanline->actual_size = tmp_size;
308         }
309     }
310
311     return 0;
312 }
313
314 static int encode_scanline_zip(EXRContext *s, const AVFrame *frame)
315 {
316     const int64_t element_size = s->pixel_type == EXR_HALF ? 2LL : 4LL;
317
318     for (int y = 0; y < s->nb_scanlines; y++) {
319         EXRScanlineData *scanline = &s->scanline[y];
320         const int scanline_height = FFMIN(s->scanline_height, frame->height - y * s->scanline_height);
321         int64_t tmp_size = element_size * s->planes * frame->width * scanline_height;
322         int64_t max_compressed_size = tmp_size * 3 / 2;
323         unsigned long actual_size, source_size;
324
325         av_fast_padded_malloc(&scanline->uncompressed_data, &scanline->uncompressed_size, tmp_size);
326         if (!scanline->uncompressed_data)
327             return AVERROR(ENOMEM);
328
329         av_fast_padded_malloc(&scanline->tmp, &scanline->tmp_size, tmp_size);
330         if (!scanline->tmp)
331             return AVERROR(ENOMEM);
332
333         av_fast_padded_malloc(&scanline->compressed_data, &scanline->compressed_size, max_compressed_size);
334         if (!scanline->compressed_data)
335             return AVERROR(ENOMEM);
336
337         switch (s->pixel_type) {
338         case EXR_FLOAT:
339             for (int l = 0; l < scanline_height; l++) {
340                 const int scanline_size = frame->width * 4 * s->planes;
341
342                 for (int p = 0; p < s->planes; p++) {
343                     int ch = s->ch_order[p];
344
345                     memcpy(scanline->uncompressed_data + scanline_size * l + p * frame->width * 4,
346                            frame->data[ch] + (y * s->scanline_height + l) * frame->linesize[ch],
347                            frame->width * 4);
348                 }
349             }
350             break;
351         case EXR_HALF:
352             for (int l = 0; l < scanline_height; l++) {
353                 const int scanline_size = frame->width * 2 * s->planes;
354
355                 for (int p = 0; p < s->planes; p++) {
356                     int ch = s->ch_order[p];
357                     uint16_t *dst = (uint16_t *)(scanline->uncompressed_data + scanline_size * l + p * frame->width * 2);
358                     uint32_t *src = (uint32_t *)(frame->data[ch] + (y * s->scanline_height + l) * frame->linesize[ch]);
359
360                     for (int x = 0; x < frame->width; x++)
361                         dst[x] = float2half(s, src[x]);
362                 }
363             }
364             break;
365         }
366
367         reorder_pixels(scanline->tmp, scanline->uncompressed_data, tmp_size);
368         predictor(scanline->tmp, tmp_size);
369         source_size = tmp_size;
370         actual_size = max_compressed_size;
371         compress(scanline->compressed_data, &actual_size,
372                  scanline->tmp, source_size);
373
374         scanline->actual_size = actual_size;
375         if (scanline->actual_size >= tmp_size) {
376             FFSWAP(uint8_t *, scanline->uncompressed_data, scanline->compressed_data);
377             FFSWAP(int, scanline->uncompressed_size, scanline->compressed_size);
378             scanline->actual_size = tmp_size;
379         }
380     }
381
382     return 0;
383 }
384
385 static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
386                         const AVFrame *frame, int *got_packet)
387 {
388     EXRContext *s = avctx->priv_data;
389     PutByteContext *pb = &s->pb;
390     int64_t offset;
391     int ret;
392     int64_t out_size = 2048LL + avctx->height * 16LL +
393                       av_image_get_buffer_size(avctx->pix_fmt,
394                                                avctx->width,
395                                                avctx->height, 64) * 3LL / 2;
396
397     if ((ret = ff_alloc_packet2(avctx, pkt, out_size, out_size)) < 0)
398         return ret;
399
400     bytestream2_init_writer(pb, pkt->data, pkt->size);
401
402     bytestream2_put_le32(pb, 20000630);
403     bytestream2_put_byte(pb, 2);
404     bytestream2_put_le24(pb, 0);
405     bytestream2_put_buffer(pb, "channels\0chlist\0", 16);
406     bytestream2_put_le32(pb, s->planes * 18 + 1);
407
408     for (int p = 0; p < s->planes; p++) {
409         bytestream2_put_byte(pb, s->ch_names[p]);
410         bytestream2_put_byte(pb, 0);
411         bytestream2_put_le32(pb, s->pixel_type);
412         bytestream2_put_le32(pb, 0);
413         bytestream2_put_le32(pb, 1);
414         bytestream2_put_le32(pb, 1);
415     }
416     bytestream2_put_byte(pb, 0);
417
418     bytestream2_put_buffer(pb, "compression\0compression\0", 24);
419     bytestream2_put_le32(pb, 1);
420     bytestream2_put_byte(pb, s->compression);
421
422     bytestream2_put_buffer(pb, "dataWindow\0box2i\0", 17);
423     bytestream2_put_le32(pb, 16);
424     bytestream2_put_le32(pb, 0);
425     bytestream2_put_le32(pb, 0);
426     bytestream2_put_le32(pb, avctx->width - 1);
427     bytestream2_put_le32(pb, avctx->height - 1);
428
429     bytestream2_put_buffer(pb, "displayWindow\0box2i\0", 20);
430     bytestream2_put_le32(pb, 16);
431     bytestream2_put_le32(pb, 0);
432     bytestream2_put_le32(pb, 0);
433     bytestream2_put_le32(pb, avctx->width - 1);
434     bytestream2_put_le32(pb, avctx->height - 1);
435
436     bytestream2_put_buffer(pb, "lineOrder\0lineOrder\0", 20);
437     bytestream2_put_le32(pb, 1);
438     bytestream2_put_byte(pb, 0);
439
440     bytestream2_put_buffer(pb, "screenWindowCenter\0v2f\0", 23);
441     bytestream2_put_le32(pb, 8);
442     bytestream2_put_le64(pb, 0);
443
444     bytestream2_put_buffer(pb, "screenWindowWidth\0float\0", 24);
445     bytestream2_put_le32(pb, 4);
446     bytestream2_put_le32(pb, av_float2int(1.f));
447
448     if (avctx->sample_aspect_ratio.num && avctx->sample_aspect_ratio.den) {
449         bytestream2_put_buffer(pb, "pixelAspectRatio\0float\0", 23);
450         bytestream2_put_le32(pb, 4);
451         bytestream2_put_le32(pb, av_float2int(av_q2d(avctx->sample_aspect_ratio)));
452     }
453
454     if (avctx->framerate.num && avctx->framerate.den) {
455         bytestream2_put_buffer(pb, "framesPerSecond\0rational\0", 25);
456         bytestream2_put_le32(pb, 8);
457         bytestream2_put_le32(pb, avctx->framerate.num);
458         bytestream2_put_le32(pb, avctx->framerate.den);
459     }
460
461     bytestream2_put_buffer(pb, "gamma\0float\0", 12);
462     bytestream2_put_le32(pb, 4);
463     bytestream2_put_le32(pb, av_float2int(s->gamma));
464
465     bytestream2_put_buffer(pb, "writer\0string\0", 14);
466     bytestream2_put_le32(pb, 4);
467     bytestream2_put_buffer(pb, "lavc", 4);
468     bytestream2_put_byte(pb, 0);
469
470     switch (s->compression) {
471     case EXR_RAW:
472         /* nothing to do */
473         break;
474     case EXR_RLE:
475         encode_scanline_rle(s, frame);
476         break;
477     case EXR_ZIP16:
478     case EXR_ZIP1:
479         encode_scanline_zip(s, frame);
480         break;
481     default:
482         av_assert0(0);
483     }
484
485     switch (s->compression) {
486     case EXR_RAW:
487         offset = bytestream2_tell_p(pb) + avctx->height * 8LL;
488
489         if (s->pixel_type == EXR_FLOAT) {
490
491             for (int y = 0; y < avctx->height; y++) {
492                 bytestream2_put_le64(pb, offset);
493                 offset += avctx->width * s->planes * 4 + 8;
494             }
495
496             for (int y = 0; y < avctx->height; y++) {
497                 bytestream2_put_le32(pb, y);
498                 bytestream2_put_le32(pb, s->planes * avctx->width * 4);
499                 for (int p = 0; p < s->planes; p++) {
500                     int ch = s->ch_order[p];
501                     bytestream2_put_buffer(pb, frame->data[ch] + y * frame->linesize[ch],
502                                            avctx->width * 4);
503                 }
504             }
505         } else {
506             for (int y = 0; y < avctx->height; y++) {
507                 bytestream2_put_le64(pb, offset);
508                 offset += avctx->width * s->planes * 2 + 8;
509             }
510
511             for (int y = 0; y < avctx->height; y++) {
512                 bytestream2_put_le32(pb, y);
513                 bytestream2_put_le32(pb, s->planes * avctx->width * 2);
514                 for (int p = 0; p < s->planes; p++) {
515                     int ch = s->ch_order[p];
516                     uint32_t *src = (uint32_t *)(frame->data[ch] + y * frame->linesize[ch]);
517
518                     for (int x = 0; x < frame->width; x++)
519                         bytestream2_put_le16(pb, float2half(s, src[x]));
520                 }
521             }
522         }
523         break;
524     case EXR_ZIP16:
525     case EXR_ZIP1:
526     case EXR_RLE:
527         offset = bytestream2_tell_p(pb) + s->nb_scanlines * 8LL;
528
529         for (int y = 0; y < s->nb_scanlines; y++) {
530             EXRScanlineData *scanline = &s->scanline[y];
531
532             bytestream2_put_le64(pb, offset);
533             offset += scanline->actual_size + 8;
534         }
535
536         for (int y = 0; y < s->nb_scanlines; y++) {
537             EXRScanlineData *scanline = &s->scanline[y];
538
539             bytestream2_put_le32(pb, y * s->scanline_height);
540             bytestream2_put_le32(pb, scanline->actual_size);
541             bytestream2_put_buffer(pb, scanline->compressed_data,
542                                    scanline->actual_size);
543         }
544         break;
545     default:
546         av_assert0(0);
547     }
548
549     av_shrink_packet(pkt, bytestream2_tell_p(pb));
550
551     pkt->flags |= AV_PKT_FLAG_KEY;
552     *got_packet = 1;
553
554     return 0;
555 }
556
557 #define OFFSET(x) offsetof(EXRContext, x)
558 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
559 static const AVOption options[] = {
560     { "compression", "set compression type", OFFSET(compression), AV_OPT_TYPE_INT,   {.i64=0}, 0, EXR_NBCOMPR-1, VE, "compr" },
561     { "none",        "none",                 0,                   AV_OPT_TYPE_CONST, {.i64=EXR_RAW}, 0, 0, VE, "compr" },
562     { "rle" ,        "RLE",                  0,                   AV_OPT_TYPE_CONST, {.i64=EXR_RLE}, 0, 0, VE, "compr" },
563     { "zip1",        "ZIP1",                 0,                   AV_OPT_TYPE_CONST, {.i64=EXR_ZIP1}, 0, 0, VE, "compr" },
564     { "zip16",       "ZIP16",                0,                   AV_OPT_TYPE_CONST, {.i64=EXR_ZIP16}, 0, 0, VE, "compr" },
565     { "format", "set pixel type", OFFSET(pixel_type), AV_OPT_TYPE_INT,   {.i64=EXR_FLOAT}, EXR_HALF, EXR_UNKNOWN-1, VE, "pixel" },
566     { "half" ,       NULL,                   0,                   AV_OPT_TYPE_CONST, {.i64=EXR_HALF},  0, 0, VE, "pixel" },
567     { "float",       NULL,                   0,                   AV_OPT_TYPE_CONST, {.i64=EXR_FLOAT}, 0, 0, VE, "pixel" },
568     { "gamma", "set gamma", OFFSET(gamma), AV_OPT_TYPE_FLOAT, {.dbl=1.f}, 0.001, FLT_MAX, VE },
569     { NULL},
570 };
571
572 static const AVClass exr_class = {
573     .class_name = "exr",
574     .item_name  = av_default_item_name,
575     .option     = options,
576     .version    = LIBAVUTIL_VERSION_INT,
577 };
578
579 AVCodec ff_exr_encoder = {
580     .name           = "exr",
581     .long_name      = NULL_IF_CONFIG_SMALL("OpenEXR image"),
582     .priv_data_size = sizeof(EXRContext),
583     .priv_class     = &exr_class,
584     .type           = AVMEDIA_TYPE_VIDEO,
585     .id             = AV_CODEC_ID_EXR,
586     .init           = encode_init,
587     .encode2        = encode_frame,
588     .close          = encode_close,
589     .capabilities   = AV_CODEC_CAP_FRAME_THREADS,
590     .pix_fmts       = (const enum AVPixelFormat[]) {
591                                                  AV_PIX_FMT_GBRPF32,
592                                                  AV_PIX_FMT_GBRAPF32,
593                                                  AV_PIX_FMT_NONE },
594 };