ffmpeg/libavcodec/exrenc.c

561 lines
18 KiB
C

/*
* Copyright (c) 2021 Paul B Mahol
*
* This file is part of FFmpeg.
*
* FFmpeg is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* FFmpeg is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with FFmpeg; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
/**
* @file
* OpenEXR encoder
*/
#include <float.h>
#include <zlib.h>
#include "libavutil/avassert.h"
#include "libavutil/opt.h"
#include "libavutil/intreadwrite.h"
#include "libavutil/imgutils.h"
#include "libavutil/pixdesc.h"
#include "avcodec.h"
#include "bytestream.h"
#include "codec_internal.h"
#include "encode.h"
#include "float2half.h"
enum ExrCompr {
EXR_RAW,
EXR_RLE,
EXR_ZIP1,
EXR_ZIP16,
EXR_NBCOMPR,
};
enum ExrPixelType {
EXR_UINT,
EXR_HALF,
EXR_FLOAT,
EXR_UNKNOWN,
};
static const char abgr_chlist[4] = { 'A', 'B', 'G', 'R' };
static const char bgr_chlist[4] = { 'B', 'G', 'R', 'A' };
static const char y_chlist[4] = { 'Y' };
static const uint8_t gbra_order[4] = { 3, 1, 0, 2 };
static const uint8_t gbr_order[4] = { 1, 0, 2, 0 };
static const uint8_t y_order[4] = { 0 };
typedef struct EXRScanlineData {
uint8_t *compressed_data;
unsigned int compressed_size;
uint8_t *uncompressed_data;
unsigned int uncompressed_size;
uint8_t *tmp;
unsigned int tmp_size;
int64_t actual_size;
} EXRScanlineData;
typedef struct EXRContext {
const AVClass *class;
int compression;
int pixel_type;
int planes;
int nb_scanlines;
int scanline_height;
float gamma;
const char *ch_names;
const uint8_t *ch_order;
PutByteContext pb;
EXRScanlineData *scanline;
uint16_t basetable[512];
uint8_t shifttable[512];
} EXRContext;
static av_cold int encode_init(AVCodecContext *avctx)
{
EXRContext *s = avctx->priv_data;
float2half_tables(s->basetable, s->shifttable);
switch (avctx->pix_fmt) {
case AV_PIX_FMT_GBRPF32:
s->planes = 3;
s->ch_names = bgr_chlist;
s->ch_order = gbr_order;
break;
case AV_PIX_FMT_GBRAPF32:
s->planes = 4;
s->ch_names = abgr_chlist;
s->ch_order = gbra_order;
break;
case AV_PIX_FMT_GRAYF32:
s->planes = 1;
s->ch_names = y_chlist;
s->ch_order = y_order;
break;
default:
av_assert0(0);
}
switch (s->compression) {
case EXR_RAW:
case EXR_RLE:
case EXR_ZIP1:
s->scanline_height = 1;
s->nb_scanlines = avctx->height;
break;
case EXR_ZIP16:
s->scanline_height = 16;
s->nb_scanlines = (avctx->height + s->scanline_height - 1) / s->scanline_height;
break;
default:
av_assert0(0);
}
s->scanline = av_calloc(s->nb_scanlines, sizeof(*s->scanline));
if (!s->scanline)
return AVERROR(ENOMEM);
return 0;
}
static av_cold int encode_close(AVCodecContext *avctx)
{
EXRContext *s = avctx->priv_data;
for (int y = 0; y < s->nb_scanlines && s->scanline; y++) {
EXRScanlineData *scanline = &s->scanline[y];
av_freep(&scanline->tmp);
av_freep(&scanline->compressed_data);
av_freep(&scanline->uncompressed_data);
}
av_freep(&s->scanline);
return 0;
}
static void reorder_pixels(uint8_t *dst, const uint8_t *src, ptrdiff_t size)
{
const ptrdiff_t half_size = (size + 1) / 2;
uint8_t *t1 = dst;
uint8_t *t2 = dst + half_size;
for (ptrdiff_t i = 0; i < half_size; i++) {
t1[i] = *(src++);
t2[i] = *(src++);
}
}
static void predictor(uint8_t *src, ptrdiff_t size)
{
int p = src[0];
for (ptrdiff_t i = 1; i < size; i++) {
int d = src[i] - p + 384;
p = src[i];
src[i] = d;
}
}
static int64_t rle_compress(uint8_t *out, int64_t out_size,
const uint8_t *in, int64_t in_size)
{
int64_t i = 0, o = 0, run = 1, copy = 0;
while (i < in_size) {
while (i + run < in_size && in[i] == in[i + run] && run < 128)
run++;
if (run >= 3) {
if (o + 2 >= out_size)
return -1;
out[o++] = run - 1;
out[o++] = in[i];
i += run;
} else {
if (i + run < in_size)
copy += run;
while (i + copy < in_size && copy < 127 && in[i + copy] != in[i + copy - 1])
copy++;
if (o + 1 + copy >= out_size)
return -1;
out[o++] = -copy;
for (int x = 0; x < copy; x++)
out[o + x] = in[i + x];
o += copy;
i += copy;
copy = 0;
}
run = 1;
}
return o;
}
static int encode_scanline_rle(EXRContext *s, const AVFrame *frame)
{
const int64_t element_size = s->pixel_type == EXR_HALF ? 2LL : 4LL;
for (int y = 0; y < frame->height; y++) {
EXRScanlineData *scanline = &s->scanline[y];
int64_t tmp_size = element_size * s->planes * frame->width;
int64_t max_compressed_size = tmp_size * 3 / 2;
av_fast_padded_malloc(&scanline->uncompressed_data, &scanline->uncompressed_size, tmp_size);
if (!scanline->uncompressed_data)
return AVERROR(ENOMEM);
av_fast_padded_malloc(&scanline->tmp, &scanline->tmp_size, tmp_size);
if (!scanline->tmp)
return AVERROR(ENOMEM);
av_fast_padded_malloc(&scanline->compressed_data, &scanline->compressed_size, max_compressed_size);
if (!scanline->compressed_data)
return AVERROR(ENOMEM);
switch (s->pixel_type) {
case EXR_FLOAT:
for (int p = 0; p < s->planes; p++) {
int ch = s->ch_order[p];
memcpy(scanline->uncompressed_data + frame->width * 4 * p,
frame->data[ch] + y * frame->linesize[ch], frame->width * 4);
}
break;
case EXR_HALF:
for (int p = 0; p < s->planes; p++) {
int ch = s->ch_order[p];
uint16_t *dst = (uint16_t *)(scanline->uncompressed_data + frame->width * 2 * p);
uint32_t *src = (uint32_t *)(frame->data[ch] + y * frame->linesize[ch]);
for (int x = 0; x < frame->width; x++)
dst[x] = float2half(src[x], s->basetable, s->shifttable);
}
break;
}
reorder_pixels(scanline->tmp, scanline->uncompressed_data, tmp_size);
predictor(scanline->tmp, tmp_size);
scanline->actual_size = rle_compress(scanline->compressed_data,
max_compressed_size,
scanline->tmp, tmp_size);
if (scanline->actual_size <= 0 || scanline->actual_size >= tmp_size) {
FFSWAP(uint8_t *, scanline->uncompressed_data, scanline->compressed_data);
FFSWAP(int, scanline->uncompressed_size, scanline->compressed_size);
scanline->actual_size = tmp_size;
}
}
return 0;
}
static int encode_scanline_zip(EXRContext *s, const AVFrame *frame)
{
const int64_t element_size = s->pixel_type == EXR_HALF ? 2LL : 4LL;
for (int y = 0; y < s->nb_scanlines; y++) {
EXRScanlineData *scanline = &s->scanline[y];
const int scanline_height = FFMIN(s->scanline_height, frame->height - y * s->scanline_height);
int64_t tmp_size = element_size * s->planes * frame->width * scanline_height;
int64_t max_compressed_size = tmp_size * 3 / 2;
unsigned long actual_size, source_size;
av_fast_padded_malloc(&scanline->uncompressed_data, &scanline->uncompressed_size, tmp_size);
if (!scanline->uncompressed_data)
return AVERROR(ENOMEM);
av_fast_padded_malloc(&scanline->tmp, &scanline->tmp_size, tmp_size);
if (!scanline->tmp)
return AVERROR(ENOMEM);
av_fast_padded_malloc(&scanline->compressed_data, &scanline->compressed_size, max_compressed_size);
if (!scanline->compressed_data)
return AVERROR(ENOMEM);
switch (s->pixel_type) {
case EXR_FLOAT:
for (int l = 0; l < scanline_height; l++) {
const int scanline_size = frame->width * 4 * s->planes;
for (int p = 0; p < s->planes; p++) {
int ch = s->ch_order[p];
memcpy(scanline->uncompressed_data + scanline_size * l + p * frame->width * 4,
frame->data[ch] + (y * s->scanline_height + l) * frame->linesize[ch],
frame->width * 4);
}
}
break;
case EXR_HALF:
for (int l = 0; l < scanline_height; l++) {
const int scanline_size = frame->width * 2 * s->planes;
for (int p = 0; p < s->planes; p++) {
int ch = s->ch_order[p];
uint16_t *dst = (uint16_t *)(scanline->uncompressed_data + scanline_size * l + p * frame->width * 2);
uint32_t *src = (uint32_t *)(frame->data[ch] + (y * s->scanline_height + l) * frame->linesize[ch]);
for (int x = 0; x < frame->width; x++)
dst[x] = float2half(src[x], s->basetable, s->shifttable);
}
}
break;
}
reorder_pixels(scanline->tmp, scanline->uncompressed_data, tmp_size);
predictor(scanline->tmp, tmp_size);
source_size = tmp_size;
actual_size = max_compressed_size;
compress(scanline->compressed_data, &actual_size,
scanline->tmp, source_size);
scanline->actual_size = actual_size;
if (scanline->actual_size >= tmp_size) {
FFSWAP(uint8_t *, scanline->uncompressed_data, scanline->compressed_data);
FFSWAP(int, scanline->uncompressed_size, scanline->compressed_size);
scanline->actual_size = tmp_size;
}
}
return 0;
}
static int encode_frame(AVCodecContext *avctx, AVPacket *pkt,
const AVFrame *frame, int *got_packet)
{
EXRContext *s = avctx->priv_data;
PutByteContext *pb = &s->pb;
int64_t offset;
int ret;
int64_t out_size = 2048LL + avctx->height * 16LL +
av_image_get_buffer_size(avctx->pix_fmt,
avctx->width,
avctx->height, 64) * 3LL / 2;
if ((ret = ff_get_encode_buffer(avctx, pkt, out_size, 0)) < 0)
return ret;
bytestream2_init_writer(pb, pkt->data, pkt->size);
bytestream2_put_le32(pb, 20000630);
bytestream2_put_byte(pb, 2);
bytestream2_put_le24(pb, 0);
bytestream2_put_buffer(pb, "channels\0chlist\0", 16);
bytestream2_put_le32(pb, s->planes * 18 + 1);
for (int p = 0; p < s->planes; p++) {
bytestream2_put_byte(pb, s->ch_names[p]);
bytestream2_put_byte(pb, 0);
bytestream2_put_le32(pb, s->pixel_type);
bytestream2_put_le32(pb, 0);
bytestream2_put_le32(pb, 1);
bytestream2_put_le32(pb, 1);
}
bytestream2_put_byte(pb, 0);
bytestream2_put_buffer(pb, "compression\0compression\0", 24);
bytestream2_put_le32(pb, 1);
bytestream2_put_byte(pb, s->compression);
bytestream2_put_buffer(pb, "dataWindow\0box2i\0", 17);
bytestream2_put_le32(pb, 16);
bytestream2_put_le32(pb, 0);
bytestream2_put_le32(pb, 0);
bytestream2_put_le32(pb, avctx->width - 1);
bytestream2_put_le32(pb, avctx->height - 1);
bytestream2_put_buffer(pb, "displayWindow\0box2i\0", 20);
bytestream2_put_le32(pb, 16);
bytestream2_put_le32(pb, 0);
bytestream2_put_le32(pb, 0);
bytestream2_put_le32(pb, avctx->width - 1);
bytestream2_put_le32(pb, avctx->height - 1);
bytestream2_put_buffer(pb, "lineOrder\0lineOrder\0", 20);
bytestream2_put_le32(pb, 1);
bytestream2_put_byte(pb, 0);
bytestream2_put_buffer(pb, "screenWindowCenter\0v2f\0", 23);
bytestream2_put_le32(pb, 8);
bytestream2_put_le64(pb, 0);
bytestream2_put_buffer(pb, "screenWindowWidth\0float\0", 24);
bytestream2_put_le32(pb, 4);
bytestream2_put_le32(pb, av_float2int(1.f));
if (avctx->sample_aspect_ratio.num && avctx->sample_aspect_ratio.den) {
bytestream2_put_buffer(pb, "pixelAspectRatio\0float\0", 23);
bytestream2_put_le32(pb, 4);
bytestream2_put_le32(pb, av_float2int(av_q2d(avctx->sample_aspect_ratio)));
}
if (avctx->framerate.num && avctx->framerate.den) {
bytestream2_put_buffer(pb, "framesPerSecond\0rational\0", 25);
bytestream2_put_le32(pb, 8);
bytestream2_put_le32(pb, avctx->framerate.num);
bytestream2_put_le32(pb, avctx->framerate.den);
}
bytestream2_put_buffer(pb, "gamma\0float\0", 12);
bytestream2_put_le32(pb, 4);
bytestream2_put_le32(pb, av_float2int(s->gamma));
bytestream2_put_buffer(pb, "writer\0string\0", 14);
bytestream2_put_le32(pb, 4);
bytestream2_put_buffer(pb, "lavc", 4);
bytestream2_put_byte(pb, 0);
switch (s->compression) {
case EXR_RAW:
/* nothing to do */
break;
case EXR_RLE:
encode_scanline_rle(s, frame);
break;
case EXR_ZIP16:
case EXR_ZIP1:
encode_scanline_zip(s, frame);
break;
default:
av_assert0(0);
}
switch (s->compression) {
case EXR_RAW:
offset = bytestream2_tell_p(pb) + avctx->height * 8LL;
if (s->pixel_type == EXR_FLOAT) {
for (int y = 0; y < avctx->height; y++) {
bytestream2_put_le64(pb, offset);
offset += avctx->width * s->planes * 4 + 8;
}
for (int y = 0; y < avctx->height; y++) {
bytestream2_put_le32(pb, y);
bytestream2_put_le32(pb, s->planes * avctx->width * 4);
for (int p = 0; p < s->planes; p++) {
int ch = s->ch_order[p];
bytestream2_put_buffer(pb, frame->data[ch] + y * frame->linesize[ch],
avctx->width * 4);
}
}
} else {
for (int y = 0; y < avctx->height; y++) {
bytestream2_put_le64(pb, offset);
offset += avctx->width * s->planes * 2 + 8;
}
for (int y = 0; y < avctx->height; y++) {
bytestream2_put_le32(pb, y);
bytestream2_put_le32(pb, s->planes * avctx->width * 2);
for (int p = 0; p < s->planes; p++) {
int ch = s->ch_order[p];
uint32_t *src = (uint32_t *)(frame->data[ch] + y * frame->linesize[ch]);
for (int x = 0; x < frame->width; x++)
bytestream2_put_le16(pb, float2half(src[x], s->basetable, s->shifttable));
}
}
}
break;
case EXR_ZIP16:
case EXR_ZIP1:
case EXR_RLE:
offset = bytestream2_tell_p(pb) + s->nb_scanlines * 8LL;
for (int y = 0; y < s->nb_scanlines; y++) {
EXRScanlineData *scanline = &s->scanline[y];
bytestream2_put_le64(pb, offset);
offset += scanline->actual_size + 8;
}
for (int y = 0; y < s->nb_scanlines; y++) {
EXRScanlineData *scanline = &s->scanline[y];
bytestream2_put_le32(pb, y * s->scanline_height);
bytestream2_put_le32(pb, scanline->actual_size);
bytestream2_put_buffer(pb, scanline->compressed_data,
scanline->actual_size);
}
break;
default:
av_assert0(0);
}
av_shrink_packet(pkt, bytestream2_tell_p(pb));
*got_packet = 1;
return 0;
}
#define OFFSET(x) offsetof(EXRContext, x)
#define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
static const AVOption options[] = {
{ "compression", "set compression type", OFFSET(compression), AV_OPT_TYPE_INT, {.i64=0}, 0, EXR_NBCOMPR-1, VE, "compr" },
{ "none", "none", 0, AV_OPT_TYPE_CONST, {.i64=EXR_RAW}, 0, 0, VE, "compr" },
{ "rle" , "RLE", 0, AV_OPT_TYPE_CONST, {.i64=EXR_RLE}, 0, 0, VE, "compr" },
{ "zip1", "ZIP1", 0, AV_OPT_TYPE_CONST, {.i64=EXR_ZIP1}, 0, 0, VE, "compr" },
{ "zip16", "ZIP16", 0, AV_OPT_TYPE_CONST, {.i64=EXR_ZIP16}, 0, 0, VE, "compr" },
{ "format", "set pixel type", OFFSET(pixel_type), AV_OPT_TYPE_INT, {.i64=EXR_FLOAT}, EXR_HALF, EXR_UNKNOWN-1, VE, "pixel" },
{ "half" , NULL, 0, AV_OPT_TYPE_CONST, {.i64=EXR_HALF}, 0, 0, VE, "pixel" },
{ "float", NULL, 0, AV_OPT_TYPE_CONST, {.i64=EXR_FLOAT}, 0, 0, VE, "pixel" },
{ "gamma", "set gamma", OFFSET(gamma), AV_OPT_TYPE_FLOAT, {.dbl=1.f}, 0.001, FLT_MAX, VE },
{ NULL},
};
static const AVClass exr_class = {
.class_name = "exr",
.item_name = av_default_item_name,
.option = options,
.version = LIBAVUTIL_VERSION_INT,
};
const FFCodec ff_exr_encoder = {
.p.name = "exr",
.p.long_name = NULL_IF_CONFIG_SMALL("OpenEXR image"),
.priv_data_size = sizeof(EXRContext),
.p.priv_class = &exr_class,
.p.type = AVMEDIA_TYPE_VIDEO,
.p.id = AV_CODEC_ID_EXR,
.p.capabilities = AV_CODEC_CAP_DR1 | AV_CODEC_CAP_FRAME_THREADS,
.init = encode_init,
FF_CODEC_ENCODE_CB(encode_frame),
.close = encode_close,
.p.pix_fmts = (const enum AVPixelFormat[]) {
AV_PIX_FMT_GRAYF32,
AV_PIX_FMT_GBRPF32,
AV_PIX_FMT_GBRAPF32,
AV_PIX_FMT_NONE },
};