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