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