• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 "codec_internal.h"
37 #include "encode.h"
38 #include "float2half.h"
39 
40 enum ExrCompr {
41     EXR_RAW,
42     EXR_RLE,
43     EXR_ZIP1,
44     EXR_ZIP16,
45     EXR_NBCOMPR,
46 };
47 
48 enum ExrPixelType {
49     EXR_UINT,
50     EXR_HALF,
51     EXR_FLOAT,
52     EXR_UNKNOWN,
53 };
54 
55 static const char abgr_chlist[4] = { 'A', 'B', 'G', 'R' };
56 static const char bgr_chlist[4] = { 'B', 'G', 'R', 'A' };
57 static const char y_chlist[4] = { 'Y' };
58 static const uint8_t gbra_order[4] = { 3, 1, 0, 2 };
59 static const uint8_t gbr_order[4] = { 1, 0, 2, 0 };
60 static const uint8_t y_order[4] = { 0 };
61 
62 typedef struct EXRScanlineData {
63     uint8_t *compressed_data;
64     unsigned int compressed_size;
65 
66     uint8_t *uncompressed_data;
67     unsigned int uncompressed_size;
68 
69     uint8_t *tmp;
70     unsigned int tmp_size;
71 
72     int64_t actual_size;
73 } EXRScanlineData;
74 
75 typedef struct EXRContext {
76     const AVClass *class;
77 
78     int compression;
79     int pixel_type;
80     int planes;
81     int nb_scanlines;
82     int scanline_height;
83     float gamma;
84     const char *ch_names;
85     const uint8_t *ch_order;
86     PutByteContext pb;
87 
88     EXRScanlineData *scanline;
89 
90     uint16_t basetable[512];
91     uint8_t shifttable[512];
92 } EXRContext;
93 
encode_init(AVCodecContext * avctx)94 static av_cold int encode_init(AVCodecContext *avctx)
95 {
96     EXRContext *s = avctx->priv_data;
97 
98     float2half_tables(s->basetable, s->shifttable);
99 
100     switch (avctx->pix_fmt) {
101     case AV_PIX_FMT_GBRPF32:
102         s->planes = 3;
103         s->ch_names = bgr_chlist;
104         s->ch_order = gbr_order;
105         break;
106     case AV_PIX_FMT_GBRAPF32:
107         s->planes = 4;
108         s->ch_names = abgr_chlist;
109         s->ch_order = gbra_order;
110         break;
111     case AV_PIX_FMT_GRAYF32:
112         s->planes = 1;
113         s->ch_names = y_chlist;
114         s->ch_order = y_order;
115         break;
116     default:
117         av_assert0(0);
118     }
119 
120     switch (s->compression) {
121     case EXR_RAW:
122     case EXR_RLE:
123     case EXR_ZIP1:
124         s->scanline_height = 1;
125         s->nb_scanlines = avctx->height;
126         break;
127     case EXR_ZIP16:
128         s->scanline_height = 16;
129         s->nb_scanlines = (avctx->height + s->scanline_height - 1) / s->scanline_height;
130         break;
131     default:
132         av_assert0(0);
133     }
134 
135     s->scanline = av_calloc(s->nb_scanlines, sizeof(*s->scanline));
136     if (!s->scanline)
137         return AVERROR(ENOMEM);
138 
139     return 0;
140 }
141 
encode_close(AVCodecContext * avctx)142 static av_cold int encode_close(AVCodecContext *avctx)
143 {
144     EXRContext *s = avctx->priv_data;
145 
146     for (int y = 0; y < s->nb_scanlines && s->scanline; y++) {
147         EXRScanlineData *scanline = &s->scanline[y];
148 
149         av_freep(&scanline->tmp);
150         av_freep(&scanline->compressed_data);
151         av_freep(&scanline->uncompressed_data);
152     }
153 
154     av_freep(&s->scanline);
155 
156     return 0;
157 }
158 
reorder_pixels(uint8_t * dst,const uint8_t * src,ptrdiff_t size)159 static void reorder_pixels(uint8_t *dst, const uint8_t *src, ptrdiff_t size)
160 {
161     const ptrdiff_t half_size = (size + 1) / 2;
162     uint8_t *t1 = dst;
163     uint8_t *t2 = dst + half_size;
164 
165     for (ptrdiff_t i = 0; i < half_size; i++) {
166         t1[i] = *(src++);
167         t2[i] = *(src++);
168     }
169 }
170 
predictor(uint8_t * src,ptrdiff_t size)171 static void predictor(uint8_t *src, ptrdiff_t size)
172 {
173     int p = src[0];
174 
175     for (ptrdiff_t i = 1; i < size; i++) {
176         int d = src[i] - p + 384;
177 
178         p = src[i];
179         src[i] = d;
180     }
181 }
182 
rle_compress(uint8_t * out,int64_t out_size,const uint8_t * in,int64_t in_size)183 static int64_t rle_compress(uint8_t *out, int64_t out_size,
184                             const uint8_t *in, int64_t in_size)
185 {
186     int64_t i = 0, o = 0, run = 1, copy = 0;
187 
188     while (i < in_size) {
189         while (i + run < in_size && in[i] == in[i + run] && run < 128)
190             run++;
191 
192         if (run >= 3) {
193             if (o + 2 >= out_size)
194                 return -1;
195             out[o++] = run - 1;
196             out[o++] = in[i];
197             i += run;
198         } else {
199             if (i + run < in_size)
200                 copy += run;
201             while (i + copy < in_size && copy < 127 && in[i + copy] != in[i + copy - 1])
202                 copy++;
203 
204             if (o + 1 + copy >= out_size)
205                 return -1;
206             out[o++] = -copy;
207 
208             for (int x = 0; x < copy; x++)
209                 out[o + x] = in[i + x];
210 
211             o += copy;
212             i += copy;
213             copy = 0;
214         }
215 
216         run = 1;
217     }
218 
219     return o;
220 }
221 
encode_scanline_rle(EXRContext * s,const AVFrame * frame)222 static int encode_scanline_rle(EXRContext *s, const AVFrame *frame)
223 {
224     const int64_t element_size = s->pixel_type == EXR_HALF ? 2LL : 4LL;
225 
226     for (int y = 0; y < frame->height; y++) {
227         EXRScanlineData *scanline = &s->scanline[y];
228         int64_t tmp_size = element_size * s->planes * frame->width;
229         int64_t max_compressed_size = tmp_size * 3 / 2;
230 
231         av_fast_padded_malloc(&scanline->uncompressed_data, &scanline->uncompressed_size, tmp_size);
232         if (!scanline->uncompressed_data)
233             return AVERROR(ENOMEM);
234 
235         av_fast_padded_malloc(&scanline->tmp, &scanline->tmp_size, tmp_size);
236         if (!scanline->tmp)
237             return AVERROR(ENOMEM);
238 
239         av_fast_padded_malloc(&scanline->compressed_data, &scanline->compressed_size, max_compressed_size);
240         if (!scanline->compressed_data)
241             return AVERROR(ENOMEM);
242 
243         switch (s->pixel_type) {
244         case EXR_FLOAT:
245             for (int p = 0; p < s->planes; p++) {
246                 int ch = s->ch_order[p];
247 
248                 memcpy(scanline->uncompressed_data + frame->width * 4 * p,
249                        frame->data[ch] + y * frame->linesize[ch], frame->width * 4);
250             }
251             break;
252         case EXR_HALF:
253             for (int p = 0; p < s->planes; p++) {
254                 int ch = s->ch_order[p];
255                 uint16_t *dst = (uint16_t *)(scanline->uncompressed_data + frame->width * 2 * p);
256                 uint32_t *src = (uint32_t *)(frame->data[ch] + y * frame->linesize[ch]);
257 
258                 for (int x = 0; x < frame->width; x++)
259                     dst[x] = float2half(src[x], s->basetable, s->shifttable);
260             }
261             break;
262         }
263 
264         reorder_pixels(scanline->tmp, scanline->uncompressed_data, tmp_size);
265         predictor(scanline->tmp, tmp_size);
266         scanline->actual_size = rle_compress(scanline->compressed_data,
267                                              max_compressed_size,
268                                              scanline->tmp, tmp_size);
269 
270         if (scanline->actual_size <= 0 || scanline->actual_size >= tmp_size) {
271             FFSWAP(uint8_t *, scanline->uncompressed_data, scanline->compressed_data);
272             FFSWAP(int, scanline->uncompressed_size, scanline->compressed_size);
273             scanline->actual_size = tmp_size;
274         }
275     }
276 
277     return 0;
278 }
279 
encode_scanline_zip(EXRContext * s,const AVFrame * frame)280 static int encode_scanline_zip(EXRContext *s, const AVFrame *frame)
281 {
282     const int64_t element_size = s->pixel_type == EXR_HALF ? 2LL : 4LL;
283 
284     for (int y = 0; y < s->nb_scanlines; y++) {
285         EXRScanlineData *scanline = &s->scanline[y];
286         const int scanline_height = FFMIN(s->scanline_height, frame->height - y * s->scanline_height);
287         int64_t tmp_size = element_size * s->planes * frame->width * scanline_height;
288         int64_t max_compressed_size = tmp_size * 3 / 2;
289         unsigned long actual_size, source_size;
290 
291         av_fast_padded_malloc(&scanline->uncompressed_data, &scanline->uncompressed_size, tmp_size);
292         if (!scanline->uncompressed_data)
293             return AVERROR(ENOMEM);
294 
295         av_fast_padded_malloc(&scanline->tmp, &scanline->tmp_size, tmp_size);
296         if (!scanline->tmp)
297             return AVERROR(ENOMEM);
298 
299         av_fast_padded_malloc(&scanline->compressed_data, &scanline->compressed_size, max_compressed_size);
300         if (!scanline->compressed_data)
301             return AVERROR(ENOMEM);
302 
303         switch (s->pixel_type) {
304         case EXR_FLOAT:
305             for (int l = 0; l < scanline_height; l++) {
306                 const int scanline_size = frame->width * 4 * s->planes;
307 
308                 for (int p = 0; p < s->planes; p++) {
309                     int ch = s->ch_order[p];
310 
311                     memcpy(scanline->uncompressed_data + scanline_size * l + p * frame->width * 4,
312                            frame->data[ch] + (y * s->scanline_height + l) * frame->linesize[ch],
313                            frame->width * 4);
314                 }
315             }
316             break;
317         case EXR_HALF:
318             for (int l = 0; l < scanline_height; l++) {
319                 const int scanline_size = frame->width * 2 * s->planes;
320 
321                 for (int p = 0; p < s->planes; p++) {
322                     int ch = s->ch_order[p];
323                     uint16_t *dst = (uint16_t *)(scanline->uncompressed_data + scanline_size * l + p * frame->width * 2);
324                     uint32_t *src = (uint32_t *)(frame->data[ch] + (y * s->scanline_height + l) * frame->linesize[ch]);
325 
326                     for (int x = 0; x < frame->width; x++)
327                         dst[x] = float2half(src[x], s->basetable, s->shifttable);
328                 }
329             }
330             break;
331         }
332 
333         reorder_pixels(scanline->tmp, scanline->uncompressed_data, tmp_size);
334         predictor(scanline->tmp, tmp_size);
335         source_size = tmp_size;
336         actual_size = max_compressed_size;
337         compress(scanline->compressed_data, &actual_size,
338                  scanline->tmp, source_size);
339 
340         scanline->actual_size = actual_size;
341         if (scanline->actual_size >= tmp_size) {
342             FFSWAP(uint8_t *, scanline->uncompressed_data, scanline->compressed_data);
343             FFSWAP(int, scanline->uncompressed_size, scanline->compressed_size);
344             scanline->actual_size = tmp_size;
345         }
346     }
347 
348     return 0;
349 }
350 
encode_frame(AVCodecContext * avctx,AVPacket * pkt,const AVFrame * frame,int * got_packet)351 static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
352                         const AVFrame *frame, int *got_packet)
353 {
354     EXRContext *s = avctx->priv_data;
355     PutByteContext *pb = &s->pb;
356     int64_t offset;
357     int ret;
358     int64_t out_size = 2048LL + avctx->height * 16LL +
359                       av_image_get_buffer_size(avctx->pix_fmt,
360                                                avctx->width,
361                                                avctx->height, 64) * 3LL / 2;
362 
363     if ((ret = ff_get_encode_buffer(avctx, pkt, out_size, 0)) < 0)
364         return ret;
365 
366     bytestream2_init_writer(pb, pkt->data, pkt->size);
367 
368     bytestream2_put_le32(pb, 20000630);
369     bytestream2_put_byte(pb, 2);
370     bytestream2_put_le24(pb, 0);
371     bytestream2_put_buffer(pb, "channels\0chlist\0", 16);
372     bytestream2_put_le32(pb, s->planes * 18 + 1);
373 
374     for (int p = 0; p < s->planes; p++) {
375         bytestream2_put_byte(pb, s->ch_names[p]);
376         bytestream2_put_byte(pb, 0);
377         bytestream2_put_le32(pb, s->pixel_type);
378         bytestream2_put_le32(pb, 0);
379         bytestream2_put_le32(pb, 1);
380         bytestream2_put_le32(pb, 1);
381     }
382     bytestream2_put_byte(pb, 0);
383 
384     bytestream2_put_buffer(pb, "compression\0compression\0", 24);
385     bytestream2_put_le32(pb, 1);
386     bytestream2_put_byte(pb, s->compression);
387 
388     bytestream2_put_buffer(pb, "dataWindow\0box2i\0", 17);
389     bytestream2_put_le32(pb, 16);
390     bytestream2_put_le32(pb, 0);
391     bytestream2_put_le32(pb, 0);
392     bytestream2_put_le32(pb, avctx->width - 1);
393     bytestream2_put_le32(pb, avctx->height - 1);
394 
395     bytestream2_put_buffer(pb, "displayWindow\0box2i\0", 20);
396     bytestream2_put_le32(pb, 16);
397     bytestream2_put_le32(pb, 0);
398     bytestream2_put_le32(pb, 0);
399     bytestream2_put_le32(pb, avctx->width - 1);
400     bytestream2_put_le32(pb, avctx->height - 1);
401 
402     bytestream2_put_buffer(pb, "lineOrder\0lineOrder\0", 20);
403     bytestream2_put_le32(pb, 1);
404     bytestream2_put_byte(pb, 0);
405 
406     bytestream2_put_buffer(pb, "screenWindowCenter\0v2f\0", 23);
407     bytestream2_put_le32(pb, 8);
408     bytestream2_put_le64(pb, 0);
409 
410     bytestream2_put_buffer(pb, "screenWindowWidth\0float\0", 24);
411     bytestream2_put_le32(pb, 4);
412     bytestream2_put_le32(pb, av_float2int(1.f));
413 
414     if (avctx->sample_aspect_ratio.num && avctx->sample_aspect_ratio.den) {
415         bytestream2_put_buffer(pb, "pixelAspectRatio\0float\0", 23);
416         bytestream2_put_le32(pb, 4);
417         bytestream2_put_le32(pb, av_float2int(av_q2d(avctx->sample_aspect_ratio)));
418     }
419 
420     if (avctx->framerate.num && avctx->framerate.den) {
421         bytestream2_put_buffer(pb, "framesPerSecond\0rational\0", 25);
422         bytestream2_put_le32(pb, 8);
423         bytestream2_put_le32(pb, avctx->framerate.num);
424         bytestream2_put_le32(pb, avctx->framerate.den);
425     }
426 
427     bytestream2_put_buffer(pb, "gamma\0float\0", 12);
428     bytestream2_put_le32(pb, 4);
429     bytestream2_put_le32(pb, av_float2int(s->gamma));
430 
431     bytestream2_put_buffer(pb, "writer\0string\0", 14);
432     bytestream2_put_le32(pb, 4);
433     bytestream2_put_buffer(pb, "lavc", 4);
434     bytestream2_put_byte(pb, 0);
435 
436     switch (s->compression) {
437     case EXR_RAW:
438         /* nothing to do */
439         break;
440     case EXR_RLE:
441         encode_scanline_rle(s, frame);
442         break;
443     case EXR_ZIP16:
444     case EXR_ZIP1:
445         encode_scanline_zip(s, frame);
446         break;
447     default:
448         av_assert0(0);
449     }
450 
451     switch (s->compression) {
452     case EXR_RAW:
453         offset = bytestream2_tell_p(pb) + avctx->height * 8LL;
454 
455         if (s->pixel_type == EXR_FLOAT) {
456 
457             for (int y = 0; y < avctx->height; y++) {
458                 bytestream2_put_le64(pb, offset);
459                 offset += avctx->width * s->planes * 4 + 8;
460             }
461 
462             for (int y = 0; y < avctx->height; y++) {
463                 bytestream2_put_le32(pb, y);
464                 bytestream2_put_le32(pb, s->planes * avctx->width * 4);
465                 for (int p = 0; p < s->planes; p++) {
466                     int ch = s->ch_order[p];
467                     bytestream2_put_buffer(pb, frame->data[ch] + y * frame->linesize[ch],
468                                            avctx->width * 4);
469                 }
470             }
471         } else {
472             for (int y = 0; y < avctx->height; y++) {
473                 bytestream2_put_le64(pb, offset);
474                 offset += avctx->width * s->planes * 2 + 8;
475             }
476 
477             for (int y = 0; y < avctx->height; y++) {
478                 bytestream2_put_le32(pb, y);
479                 bytestream2_put_le32(pb, s->planes * avctx->width * 2);
480                 for (int p = 0; p < s->planes; p++) {
481                     int ch = s->ch_order[p];
482                     uint32_t *src = (uint32_t *)(frame->data[ch] + y * frame->linesize[ch]);
483 
484                     for (int x = 0; x < frame->width; x++)
485                         bytestream2_put_le16(pb, float2half(src[x], s->basetable, s->shifttable));
486                 }
487             }
488         }
489         break;
490     case EXR_ZIP16:
491     case EXR_ZIP1:
492     case EXR_RLE:
493         offset = bytestream2_tell_p(pb) + s->nb_scanlines * 8LL;
494 
495         for (int y = 0; y < s->nb_scanlines; y++) {
496             EXRScanlineData *scanline = &s->scanline[y];
497 
498             bytestream2_put_le64(pb, offset);
499             offset += scanline->actual_size + 8;
500         }
501 
502         for (int y = 0; y < s->nb_scanlines; y++) {
503             EXRScanlineData *scanline = &s->scanline[y];
504 
505             bytestream2_put_le32(pb, y * s->scanline_height);
506             bytestream2_put_le32(pb, scanline->actual_size);
507             bytestream2_put_buffer(pb, scanline->compressed_data,
508                                    scanline->actual_size);
509         }
510         break;
511     default:
512         av_assert0(0);
513     }
514 
515     av_shrink_packet(pkt, bytestream2_tell_p(pb));
516 
517     *got_packet = 1;
518 
519     return 0;
520 }
521 
522 #define OFFSET(x) offsetof(EXRContext, x)
523 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
524 static const AVOption options[] = {
525     { "compression", "set compression type", OFFSET(compression), AV_OPT_TYPE_INT,   {.i64=0}, 0, EXR_NBCOMPR-1, VE, "compr" },
526     { "none",        "none",                 0,                   AV_OPT_TYPE_CONST, {.i64=EXR_RAW}, 0, 0, VE, "compr" },
527     { "rle" ,        "RLE",                  0,                   AV_OPT_TYPE_CONST, {.i64=EXR_RLE}, 0, 0, VE, "compr" },
528     { "zip1",        "ZIP1",                 0,                   AV_OPT_TYPE_CONST, {.i64=EXR_ZIP1}, 0, 0, VE, "compr" },
529     { "zip16",       "ZIP16",                0,                   AV_OPT_TYPE_CONST, {.i64=EXR_ZIP16}, 0, 0, VE, "compr" },
530     { "format", "set pixel type", OFFSET(pixel_type), AV_OPT_TYPE_INT,   {.i64=EXR_FLOAT}, EXR_HALF, EXR_UNKNOWN-1, VE, "pixel" },
531     { "half" ,       NULL,                   0,                   AV_OPT_TYPE_CONST, {.i64=EXR_HALF},  0, 0, VE, "pixel" },
532     { "float",       NULL,                   0,                   AV_OPT_TYPE_CONST, {.i64=EXR_FLOAT}, 0, 0, VE, "pixel" },
533     { "gamma", "set gamma", OFFSET(gamma), AV_OPT_TYPE_FLOAT, {.dbl=1.f}, 0.001, FLT_MAX, VE },
534     { NULL},
535 };
536 
537 static const AVClass exr_class = {
538     .class_name = "exr",
539     .item_name  = av_default_item_name,
540     .option     = options,
541     .version    = LIBAVUTIL_VERSION_INT,
542 };
543 
544 const FFCodec ff_exr_encoder = {
545     .p.name         = "exr",
546     .p.long_name    = NULL_IF_CONFIG_SMALL("OpenEXR image"),
547     .priv_data_size = sizeof(EXRContext),
548     .p.priv_class   = &exr_class,
549     .p.type         = AVMEDIA_TYPE_VIDEO,
550     .p.id           = AV_CODEC_ID_EXR,
551     .p.capabilities = AV_CODEC_CAP_DR1 | AV_CODEC_CAP_FRAME_THREADS,
552     .init           = encode_init,
553     FF_CODEC_ENCODE_CB(encode_frame),
554     .close          = encode_close,
555     .p.pix_fmts     = (const enum AVPixelFormat[]) {
556                                                  AV_PIX_FMT_GRAYF32,
557                                                  AV_PIX_FMT_GBRPF32,
558                                                  AV_PIX_FMT_GBRAPF32,
559                                                  AV_PIX_FMT_NONE },
560     .caps_internal  = FF_CODEC_CAP_INIT_THREADSAFE,
561 };
562