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