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 40enum ExrCompr { 41 EXR_RAW, 42 EXR_RLE, 43 EXR_ZIP1, 44 EXR_ZIP16, 45 EXR_NBCOMPR, 46}; 47 48enum ExrPixelType { 49 EXR_UINT, 50 EXR_HALF, 51 EXR_FLOAT, 52 EXR_UNKNOWN, 53}; 54 55static const char abgr_chlist[4] = { 'A', 'B', 'G', 'R' }; 56static const char bgr_chlist[4] = { 'B', 'G', 'R', 'A' }; 57static const char y_chlist[4] = { 'Y' }; 58static const uint8_t gbra_order[4] = { 3, 1, 0, 2 }; 59static const uint8_t gbr_order[4] = { 1, 0, 2, 0 }; 60static const uint8_t y_order[4] = { 0 }; 61 62typedef 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 75typedef 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 94static 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 142static 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 159static 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 171static 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 183static 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 222static 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 280static 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 351static 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 524static 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 537static 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 544const 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