[FFmpeg-devel] [PATCH v2] avcodec/pngenc: support writing iCCP chunks

[Niklas Haas]

From: Niklas Haas <git at haasn.dev>

encode_zbuf is mostly a mirror image of decode_zbuf. Other than that,
the code is pretty straightforward. Special care needs to be taken to
avoid writing more than 79 characters of the profile description (the
maximum supported).

Also add a FATE transcode test to ensure that the ICC profile gets
encoded correctly.
---
Oops. `-c copy` doesn't actually test the PNG writing code. Need to use
`-c png` instead. Fixed in v2.
---
 libavcodec/pngenc.c    | 77 +++++++++++++++++++++++++++++++++++++++++-
 tests/fate/image.mak   |  3 ++
 tests/ref/fate/png-icc |  8 +++++
 3 files changed, 87 insertions(+), 1 deletion(-)
 create mode 100644 tests/ref/fate/png-icc

diff --git a/libavcodec/pngenc.c b/libavcodec/pngenc.c
index 3ebcc1e571..24530bb62f 100644
--- a/libavcodec/pngenc.c
+++ b/libavcodec/pngenc.c
@@ -28,6 +28,7 @@
 #include "apng.h"
 
 #include "libavutil/avassert.h"
+#include "libavutil/bprint.h"
 #include "libavutil/crc.h"
 #include "libavutil/libm.h"
 #include "libavutil/opt.h"
@@ -343,6 +344,65 @@ static int png_get_gama(enum AVColorTransferCharacteristic trc, uint8_t *buf)
     return 1;
 }
 
+static int encode_zbuf(AVBPrint *bp, const uint8_t *data, size_t size)
+{
+    z_stream zstream;
+    unsigned char *buf;
+    unsigned buf_size;
+    int ret;
+
+    zstream.zalloc = ff_png_zalloc,
+    zstream.zfree  = ff_png_zfree,
+    zstream.opaque = NULL;
+    if (deflateInit(&zstream, Z_DEFAULT_COMPRESSION) != Z_OK)
+        return AVERROR_EXTERNAL;
+    zstream.next_in  = data;
+    zstream.avail_in = size;
+
+    for (;;) {
+        av_bprint_get_buffer(bp, 2, &buf, &buf_size);
+        if (buf_size < 2) {
+            deflateEnd(&zstream);
+            return AVERROR(ENOMEM);
+        }
+
+        zstream.next_out  = buf;
+        zstream.avail_out = buf_size - 1;
+        ret = deflate(&zstream, Z_FINISH);
+        if (ret != Z_OK && ret != Z_STREAM_END) {
+            deflateEnd(&zstream);
+            return AVERROR_EXTERNAL;
+        }
+
+        bp->len += zstream.next_out - buf;
+        if (ret == Z_STREAM_END) {
+            deflateEnd(&zstream);
+            return 0;
+        }
+    }
+}
+
+static int png_get_iccp(AVBPrint *bp, const AVFrameSideData *sd, char **buf_out)
+{
+    const AVDictionaryEntry *name;
+    int ret;
+
+    av_bprint_init(bp, 0, AV_BPRINT_SIZE_UNLIMITED);
+
+    /* profile header */
+    name = av_dict_get(sd->metadata, "name", NULL, 0);
+    av_bprintf(bp, "%.79s", (name && name->value[0]) ? name->value : "icc");
+    av_bprint_chars(bp, 0, 2); /* terminating \0 and compression method */
+
+    /* profile data */
+    if ((ret = encode_zbuf(bp, sd->data, sd->size))) {
+        av_bprint_finalize(bp, NULL);
+        return ret;
+    }
+
+    return av_bprint_finalize(bp, buf_out);
+}
+
 static int encode_headers(AVCodecContext *avctx, const AVFrame *pict)
 {
     AVFrameSideData *side_data;
@@ -399,7 +459,22 @@ static int encode_headers(AVCodecContext *avctx, const AVFrame *pict)
     if (png_get_gama(pict->color_trc, s->buf))
         png_write_chunk(&s->bytestream, MKTAG('g', 'A', 'M', 'A'), s->buf, 4);
 
-    /* put the palette if needed */
+    side_data = av_frame_get_side_data(pict, AV_FRAME_DATA_ICC_PROFILE);
+    if (side_data && side_data->size) {
+        AVBPrint bp;
+        char *buf;
+        int ret;
+
+        if ((ret = png_get_iccp(&bp, side_data, &buf))) {
+            av_log(avctx, AV_LOG_WARNING, "Failed writing iCCP chunk: %s\n",
+                   av_err2str(ret));
+        } else {
+            png_write_chunk(&s->bytestream, MKTAG('i', 'C', 'C', 'P'), buf, bp.size);
+            av_free(buf);
+        }
+    }
+
+    /* put the palette if needed, must be after colorspace information */
     if (s->color_type == PNG_COLOR_TYPE_PALETTE) {
         int has_alpha, alpha, i;
         unsigned int v;
diff --git a/tests/fate/image.mak b/tests/fate/image.mak
index 573d398915..da4f3709e9 100644
--- a/tests/fate/image.mak
+++ b/tests/fate/image.mak
@@ -385,6 +385,9 @@ FATE_PNG_PROBE += fate-png-side-data
 fate-png-side-data: CMD = run ffprobe$(PROGSSUF)$(EXESUF) -show_frames \
     -i $(TARGET_SAMPLES)/png1/lena-int_rgb24.png
 
+FATE_PNG_PROBE += fate-png-icc
+fate-png-icc: CMD = transcode png_pipe $(TARGET_SAMPLES)/png1/lena-int_rgb24.png image2 "-c png"
+
 FATE_PNG-$(call DEMDEC, IMAGE2, PNG) += $(FATE_PNG)
 FATE_PNG_PROBE-$(call DEMDEC, IMAGE2, PNG) += $(FATE_PNG_PROBE)
 FATE_IMAGE += $(FATE_PNG-yes)
diff --git a/tests/ref/fate/png-icc b/tests/ref/fate/png-icc
new file mode 100644
index 0000000000..d3cf55263e
--- /dev/null
+++ b/tests/ref/fate/png-icc
@@ -0,0 +1,8 @@
+7e412f6a9e2c7fcb674336e5c104518d *tests/data/fate/png-icc.image2
+49398 tests/data/fate/png-icc.image2
+#tb 0: 1/25
+#media_type 0: video
+#codec_id 0: rawvideo
+#dimensions 0: 128x128
+#sar 0: 2835/2835
+0,          0,          0,        1,    49152, 0xe0013dee
-- 
2.35.1