[FFmpeg-devel] avcodec/mjpegenc: support writing ICC profiles

diff --git a/libavcodec/ljpegenc.c b/libavcodec/ljpegenc.c
index e15f448f90..c54450c338 100644
--- a/libavcodec/ljpegenc.c
+++ b/libavcodec/ljpegenc.c
@@ -216,6 +216,7 @@  static int ljpeg_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
 {
     LJpegEncContext *s = avctx->priv_data;
     PutBitContext pb;
+    const AVFrameSideData *sd;
     const int width  = avctx->width;
     const int height = avctx->height;
     const int mb_width  = (width  + s->hsample[0] - 1) / s->hsample[0];
@@ -233,12 +234,15 @@  static int ljpeg_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
                         * s->hsample[0] * s->vsample[0];
     }
 
+    if ((sd = av_frame_get_side_data(pict, AV_FRAME_DATA_ICC_PROFILE)))
+        max_pkt_size += sd->size;
+
     if ((ret = ff_alloc_packet(avctx, pkt, max_pkt_size)) < 0)
         return ret;
 
     init_put_bits(&pb, pkt->data, pkt->size);
 
-    ff_mjpeg_encode_picture_header(avctx, &pb, NULL, &s->scantable,
+    ff_mjpeg_encode_picture_header(avctx, &pb, pict, NULL, &s->scantable,
                                    s->pred, s->matrix, s->matrix);
 
     header_bits = put_bits_count(&pb);
diff --git a/libavcodec/mjpegenc.c b/libavcodec/mjpegenc.c
index 08671b0df7..d7c0c763a1 100644
--- a/libavcodec/mjpegenc.c
+++ b/libavcodec/mjpegenc.c
@@ -77,7 +77,8 @@  static av_cold void init_uni_ac_vlc(const uint8_t huff_size_ac[256],
 
 static void mjpeg_encode_picture_header(MpegEncContext *s)
 {
-    ff_mjpeg_encode_picture_header(s->avctx, &s->pb, s->mjpeg_ctx,
+    const AVFrame *frame = s->picture->f;
+    ff_mjpeg_encode_picture_header(s->avctx, &s->pb, frame, s->mjpeg_ctx,
                                    &s->intra_scantable, 0,
                                    s->intra_matrix, s->chroma_intra_matrix);
 
diff --git a/libavcodec/mjpegenc_common.c b/libavcodec/mjpegenc_common.c
index 995e2b7670..099ed3fce3 100644
--- a/libavcodec/mjpegenc_common.c
+++ b/libavcodec/mjpegenc_common.c
@@ -130,8 +130,10 @@  static void jpeg_table_header(AVCodecContext *avctx, PutBitContext *p,
     AV_WB16(ptr, size);
 }
 
-static void jpeg_put_comments(AVCodecContext *avctx, PutBitContext *p)
+static void jpeg_put_comments(AVCodecContext *avctx, PutBitContext *p,
+                              const AVFrame *frame)
 {
+    const AVFrameSideData *sd;
     int size;
     uint8_t *ptr;
 
@@ -161,6 +163,39 @@  static void jpeg_put_comments(AVCodecContext *avctx, PutBitContext *p)
         put_bits(p, 8, 0); /* thumbnail height */
     }
 
+    /* ICC profile */
+    if ((sd = av_frame_get_side_data(frame, AV_FRAME_DATA_ICC_PROFILE))) {
+        static const uint8_t hdr_size = strlen("ICC_PROFILE")+5;
+        static const uint16_t max_chunk_size = UINT16_MAX - hdr_size;
+        const uint8_t *data = sd->data;
+        size_t remaining = sd->size;
+        size_t nb_chunks = (remaining + max_chunk_size - 1) / max_chunk_size;
+        if (nb_chunks > 255) {
+            av_log(avctx, AV_LOG_WARNING,
+                "Cannot store %zu byte ICC profile: too large for JPEG\n",
+                sd->size);
+            nb_chunks = remaining = 0;
+        }
+
+        for (int i = 0; i < nb_chunks; i++) {
+            size = FFMIN(remaining, max_chunk_size);
+            av_assert1(size > 0);
+            put_marker(p, APP2);
+            put_bits(p, 16, size + hdr_size);
+            ff_put_string(p, "ICC_PROFILE", 1);
+            put_bits(p, 8, i+1);
+            put_bits(p, 8, nb_chunks);
+            flush_put_bits(p);
+            ptr = put_bits_ptr(p);
+            skip_put_bytes(p, size);
+            memcpy(ptr, data, size);
+            remaining -= size;
+            data += size;
+        }
+
+        av_assert1(!remaining);
+    }
+
     /* comment */
     if (!(avctx->flags & AV_CODEC_FLAG_BITEXACT)) {
         put_marker(p, COM);
@@ -213,7 +248,7 @@  void ff_mjpeg_init_hvsample(AVCodecContext *avctx, int hsample[4], int vsample[4
 }
 
 void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb,
-                                    MJpegContext *m,
+                                    const AVFrame *frame, struct MJpegContext *m,
                                     ScanTable *intra_scantable, int pred,
                                     uint16_t luma_intra_matrix[64],
                                     uint16_t chroma_intra_matrix[64])
@@ -233,7 +268,7 @@  void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb,
     if (avctx->codec_id == AV_CODEC_ID_AMV)
         return;
 
-    jpeg_put_comments(avctx, pb);
+    jpeg_put_comments(avctx, pb, frame);
 
     jpeg_table_header(avctx, pb, m, intra_scantable,
                       luma_intra_matrix, chroma_intra_matrix, hsample);
diff --git a/libavcodec/mjpegenc_common.h b/libavcodec/mjpegenc_common.h
index ac753bf153..3c0a7addac 100644
--- a/libavcodec/mjpegenc_common.h
+++ b/libavcodec/mjpegenc_common.h
@@ -30,7 +30,7 @@ 
 struct MJpegContext;
 
 void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb,
-                                    struct MJpegContext *m,
+                                    const AVFrame *frame, struct MJpegContext *m,
                                     ScanTable *intra_scantable, int pred,
                                     uint16_t luma_intra_matrix[64],
                                     uint16_t chroma_intra_matrix[64]);
diff --git a/libavcodec/mpegvideo_enc.c b/libavcodec/mpegvideo_enc.c
index c69114ea15..932dfa34f7 100644
--- a/libavcodec/mpegvideo_enc.c
+++ b/libavcodec/mpegvideo_enc.c
@@ -1686,10 +1686,13 @@  int ff_mpv_encode_picture(AVCodecContext *avctx, AVPacket *pkt,
 
     /* output? */
     if (s->new_picture.f->data[0]) {
+        const AVFrameSideData *sd;
         int growing_buffer = context_count == 1 && !pkt->data && !s->data_partitioning;
         int pkt_size = growing_buffer ? FFMAX(s->mb_width*s->mb_height*64+10000, avctx->internal->byte_buffer_size) - AV_INPUT_BUFFER_PADDING_SIZE
                                               :
                                               s->mb_width*s->mb_height*(MAX_MB_BYTES+100)+10000;
+        if ((sd = av_frame_get_side_data(s->new_picture.f, AV_FRAME_DATA_ICC_PROFILE)))
+            pkt_size += sd->size;
         if ((ret = ff_alloc_packet(avctx, pkt, pkt_size)) < 0)
             return ret;
         if (s->mb_info) {