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
  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         unsigned long actual_size, source_size;
 255
 256         av_fast_padded_malloc(&scanline->uncompressed_data, &scanline->uncompressed_size, tmp_size);
 257         if (!scanline->uncompressed_data)
 258             return AVERROR(ENOMEM);
 259
 260         av_fast_padded_malloc(&scanline->tmp, &scanline->tmp_size, tmp_size);
 261         if (!scanline->tmp)
 262             return AVERROR(ENOMEM);
 263
 264         av_fast_padded_malloc(&scanline->compressed_data, &scanline->compressed_size, max_compressed_size);
 265         if (!scanline->compressed_data)
 266             return AVERROR(ENOMEM);
 267
 268         for (int l = 0; l < scanline_height; l++) {
 269             const int scanline_size = frame->width * 4 * s->planes;
 270
 271             for (int p = 0; p < s->planes; p++) {
 272                 int ch = s->ch_order[p];
 273
 274                 memcpy(scanline->uncompressed_data + scanline_size * l + p * frame->width * 4,
 275                        frame->data[ch] + (y * s->scanline_height + l) * frame->linesize[ch],
 276                        frame->width * 4);
 277             }
 278         }
 279
 280         reorder_pixels(scanline->tmp, scanline->uncompressed_data, tmp_size);
 281         predictor(scanline->tmp, tmp_size);
 282         source_size = tmp_size;
 283         actual_size = max_compressed_size;
 284         compress(scanline->compressed_data, &actual_size,
 285                  scanline->tmp, source_size);
 286
 287         scanline->actual_size = actual_size;
 288         if (scanline->actual_size >= tmp_size) {
 289             FFSWAP(uint8_t *, scanline->uncompressed_data, scanline->compressed_data);
 290             FFSWAP(int, scanline->uncompressed_size, scanline->compressed_size);
 291             scanline->actual_size = tmp_size;
 292         }
 293     }
 294
 295     return 0;
 296 }
 297
 298 static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
 299                         const AVFrame *frame, int *got_packet)
 300 {
 301     EXRContext *s = avctx->priv_data;
 302     PutByteContext *pb = &s->pb;
 303     int64_t offset;
 304     int ret;
 305     int64_t out_size = 2048LL + avctx->height * 16LL +
 306                       av_image_get_buffer_size(avctx->pix_fmt,
 307                                                avctx->width,
 308                                                avctx->height, 64) * 3LL / 2;
 309
 310     if ((ret = ff_alloc_packet2(avctx, pkt, out_size, out_size)) < 0)
 311         return ret;
 312
 313     bytestream2_init_writer(pb, pkt->data, pkt->size);
 314
 315     bytestream2_put_le32(pb, 20000630);
 316     bytestream2_put_byte(pb, 2);
 317     bytestream2_put_le24(pb, 0);
 318     bytestream2_put_buffer(pb, "channels\0chlist\0", 16);
 319     bytestream2_put_le32(pb, s->planes * 18 + 1);
 320
 321     for (int p = 0; p < s->planes; p++) {
 322         bytestream2_put_byte(pb, s->ch_names[p]);
 323         bytestream2_put_byte(pb, 0);
 324         bytestream2_put_le32(pb, EXR_FLOAT);
 325         bytestream2_put_le32(pb, 0);
 326         bytestream2_put_le32(pb, 1);
 327         bytestream2_put_le32(pb, 1);
 328     }
 329     bytestream2_put_byte(pb, 0);
 330
 331     bytestream2_put_buffer(pb, "compression\0compression\0", 24);
 332     bytestream2_put_le32(pb, 1);
 333     bytestream2_put_byte(pb, s->compression);
 334
 335     bytestream2_put_buffer(pb, "dataWindow\0box2i\0", 17);
 336     bytestream2_put_le32(pb, 16);
 337     bytestream2_put_le32(pb, 0);
 338     bytestream2_put_le32(pb, 0);
 339     bytestream2_put_le32(pb, avctx->width - 1);
 340     bytestream2_put_le32(pb, avctx->height - 1);
 341
 342     bytestream2_put_buffer(pb, "displayWindow\0box2i\0", 20);
 343     bytestream2_put_le32(pb, 16);
 344     bytestream2_put_le32(pb, 0);
 345     bytestream2_put_le32(pb, 0);
 346     bytestream2_put_le32(pb, avctx->width - 1);
 347     bytestream2_put_le32(pb, avctx->height - 1);
 348
 349     bytestream2_put_buffer(pb, "lineOrder\0lineOrder\0", 20);
 350     bytestream2_put_le32(pb, 1);
 351     bytestream2_put_byte(pb, 0);
 352
 353     bytestream2_put_buffer(pb, "screenWindowCenter\0v2f\0", 23);
 354     bytestream2_put_le32(pb, 8);
 355     bytestream2_put_le64(pb, 0);
 356
 357     bytestream2_put_buffer(pb, "screenWindowWidth\0float\0", 24);
 358     bytestream2_put_le32(pb, 4);
 359     bytestream2_put_le32(pb, av_float2int(1.f));
 360
 361     if (avctx->sample_aspect_ratio.num && avctx->sample_aspect_ratio.den) {
 362         bytestream2_put_buffer(pb, "pixelAspectRatio\0float\0", 23);
 363         bytestream2_put_le32(pb, 4);
 364         bytestream2_put_le32(pb, av_float2int(av_q2d(avctx->sample_aspect_ratio)));
 365     }
 366
 367     if (avctx->framerate.num && avctx->framerate.den) {
 368         bytestream2_put_buffer(pb, "framesPerSecond\0rational\0", 25);
 369         bytestream2_put_le32(pb, 8);
 370         bytestream2_put_le32(pb, avctx->framerate.num);
 371         bytestream2_put_le32(pb, avctx->framerate.den);
 372     }
 373
 374     bytestream2_put_buffer(pb, "gamma\0float\0", 12);
 375     bytestream2_put_le32(pb, 4);
 376     bytestream2_put_le32(pb, av_float2int(s->gamma));
 377
 378     bytestream2_put_buffer(pb, "writer\0string\0", 14);
 379     bytestream2_put_le32(pb, 4);
 380     bytestream2_put_buffer(pb, "lavc", 4);
 381     bytestream2_put_byte(pb, 0);
 382
 383     switch (s->compression) {
 384     case EXR_RAW:
 385         /* nothing to do */
 386         break;
 387     case EXR_RLE:
 388         encode_scanline_rle(s, frame);
 389         break;
 390     case EXR_ZIP16:
 391     case EXR_ZIP1:
 392         encode_scanline_zip(s, frame);
 393         break;
 394     default:
 395         av_assert0(0);
 396     }
 397
 398     switch (s->compression) {
 399     case EXR_RAW:
 400         offset = bytestream2_tell_p(pb) + avctx->height * 8LL;
 401
 402         for (int y = 0; y < avctx->height; y++) {
 403             bytestream2_put_le64(pb, offset);
 404             offset += avctx->width * s->planes * 4 + 8;
 405         }
 406
 407         for (int y = 0; y < avctx->height; y++) {
 408             bytestream2_put_le32(pb, y);
 409             bytestream2_put_le32(pb, s->planes * avctx->width * 4);
 410             for (int p = 0; p < s->planes; p++) {
 411                 int ch = s->ch_order[p];
 412                 bytestream2_put_buffer(pb, frame->data[ch] + y * frame->linesize[ch],
 413                                        avctx->width * 4);
 414             }
 415         }
 416         break;
 417     case EXR_ZIP16:
 418     case EXR_ZIP1:
 419     case EXR_RLE:
 420         offset = bytestream2_tell_p(pb) + s->nb_scanlines * 8LL;
 421
 422         for (int y = 0; y < s->nb_scanlines; y++) {
 423             EXRScanlineData *scanline = &s->scanline[y];
 424
 425             bytestream2_put_le64(pb, offset);
 426             offset += scanline->actual_size + 8;
 427         }
 428
 429         for (int y = 0; y < s->nb_scanlines; y++) {
 430             EXRScanlineData *scanline = &s->scanline[y];
 431
 432             bytestream2_put_le32(pb, y * s->scanline_height);
 433             bytestream2_put_le32(pb, scanline->actual_size);
 434             bytestream2_put_buffer(pb, scanline->compressed_data,
 435                                    scanline->actual_size);
 436         }
 437         break;
 438     default:
 439         av_assert0(0);
 440     }
 441
 442     av_shrink_packet(pkt, bytestream2_tell_p(pb));
 443
 444     pkt->flags |= AV_PKT_FLAG_KEY;
 445     *got_packet = 1;
 446
 447     return 0;
 448 }
 449
 450 #define OFFSET(x) offsetof(EXRContext, x)
 451 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
 452 static const AVOption options[] = {
 453     { "compression", "set compression type", OFFSET(compression), AV_OPT_TYPE_INT,   {.i64=0}, 0, EXR_NBCOMPR-1, VE, "compr" },
 454     { "none",        "none",                 0,                   AV_OPT_TYPE_CONST, {.i64=EXR_RAW}, 0, 0, VE, "compr" },
 455     { "rle" ,        "RLE",                  0,                   AV_OPT_TYPE_CONST, {.i64=EXR_RLE}, 0, 0, VE, "compr" },
 456     { "zip1",        "ZIP1",                 0,                   AV_OPT_TYPE_CONST, {.i64=EXR_ZIP1}, 0, 0, VE, "compr" },
 457     { "zip16",       "ZIP16",                0,                   AV_OPT_TYPE_CONST, {.i64=EXR_ZIP16}, 0, 0, VE, "compr" },
 458     { "gamma", "set gamma", OFFSET(gamma), AV_OPT_TYPE_FLOAT, {.dbl=1.f}, 0.001, FLT_MAX, VE },
 459     { NULL},
 460 };
 461
 462 static const AVClass exr_class = {
 463     .class_name = "exr",
 464     .item_name  = av_default_item_name,
 465     .option     = options,
 466     .version    = LIBAVUTIL_VERSION_INT,
 467 };
 468
 469 AVCodec ff_exr_encoder = {
 470     .name           = "exr",
 471     .long_name      = NULL_IF_CONFIG_SMALL("OpenEXR image"),
 472     .priv_data_size = sizeof(EXRContext),
 473     .priv_class     = &exr_class,
 474     .type           = AVMEDIA_TYPE_VIDEO,
 475     .id             = AV_CODEC_ID_EXR,
 476     .init           = encode_init,
 477     .encode2        = encode_frame,
 478     .close          = encode_close,
 479     .capabilities   = AV_CODEC_CAP_FRAME_THREADS,
 480     .pix_fmts       = (const enum AVPixelFormat[]) {
 481                                                  AV_PIX_FMT_GBRPF32,
 482                                                  AV_PIX_FMT_GBRAPF32,
 483                                                  AV_PIX_FMT_NONE },
 484 };