git.sesse.net Git - ffmpeg/blob - libavcodec/exrenc.c

   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
  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
 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
 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
 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
 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
 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
 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
 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 const 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 };