Support reference picture defined by bitmask in MJPEG's SOS decoder

Signed-off-by: Anton Khirnov <anton@khirnov.net>

3 files changed, 73 insertions, 16 deletions

diff --git a/libavcodec/mjpegbdec.c b/libavcodec/mjpegbdec.c
index 2337f77eb1..f3366ffd6f 100644
--- a/libavcodec/mjpegbdec.c
+++ b/libavcodec/mjpegbdec.c
@@ -113,7 +113,7 @@ read_header:
         init_get_bits(&s->gb, buf_ptr+sos_offs, field_size*8);
         s->mjpb_skiptosod = (sod_offs - sos_offs - show_bits(&s->gb, 16));
         s->start_code = SOS;
-        ff_mjpeg_decode_sos(s);
+        ff_mjpeg_decode_sos(s, NULL, NULL);
     }
 
     if (s->interlaced) {
diff --git a/libavcodec/mjpegdec.c b/libavcodec/mjpegdec.c
index 156563d723..a230fa9529 100644
--- a/libavcodec/mjpegdec.c
+++ b/libavcodec/mjpegdec.c
@@ -771,10 +771,32 @@ static int ljpeg_decode_yuv_scan(MJpegDecodeContext *s, int predictor, int point
     return 0;
 }
 
-static int mjpeg_decode_scan(MJpegDecodeContext *s, int nb_components, int Ah, int Al){
+static av_always_inline void mjpeg_copy_block(uint8_t *dst, const uint8_t *src,
+                                              int linesize, int lowres)
+{
+    switch (lowres) {
+    case 0: copy_block8(dst, src, linesize, linesize, 8);
+        break;
+    case 1: copy_block4(dst, src, linesize, linesize, 4);
+        break;
+    case 2: copy_block2(dst, src, linesize, linesize, 2);
+        break;
+    case 3: *dst = *src;
+        break;
+    }
+}
+
+static int mjpeg_decode_scan(MJpegDecodeContext *s, int nb_components, int Ah, int Al,
+                             const uint8_t *mb_bitmask, const AVFrame *reference){
     int i, mb_x, mb_y;
     uint8_t* data[MAX_COMPONENTS];
+    const uint8_t *reference_data[MAX_COMPONENTS];
     int linesize[MAX_COMPONENTS];
+    GetBitContext mb_bitmask_gb;
+
+    if (mb_bitmask) {
+        init_get_bits(&mb_bitmask_gb, mb_bitmask, s->mb_width*s->mb_height);
+    }
 
     if(s->flipped && s->avctx->flags & CODEC_FLAG_EMU_EDGE) {
         av_log(s->avctx, AV_LOG_ERROR, "Can not flip image with CODEC_FLAG_EMU_EDGE set!\n");
@@ -783,23 +805,29 @@ static int mjpeg_decode_scan(MJpegDecodeContext *s, int nb_components, int Ah, i
     for(i=0; i < nb_components; i++) {
         int c = s->comp_index[i];
         data[c] = s->picture.data[c];
+        reference_data[c] = reference ? reference->data[c] : NULL;
         linesize[c]=s->linesize[c];
         s->coefs_finished[c] |= 1;
         if(s->flipped) {
             //picture should be flipped upside-down for this codec
-            data[c] += (linesize[c] * (s->v_scount[i] * (8 * s->mb_height -((s->height/s->v_max)&7)) - 1 ));
+            int offset = (linesize[c] * (s->v_scount[i] * (8 * s->mb_height -((s->height/s->v_max)&7)) - 1 ));
+            data[c] += offset;
+            reference_data[c] += offset;
             linesize[c] *= -1;
         }
     }
 
     for(mb_y = 0; mb_y < s->mb_height; mb_y++) {
         for(mb_x = 0; mb_x < s->mb_width; mb_x++) {
+            const int copy_mb = mb_bitmask && !get_bits1(&mb_bitmask_gb);
+
             if (s->restart_interval && !s->restart_count)
                 s->restart_count = s->restart_interval;
 
             for(i=0;i<nb_components;i++) {
                 uint8_t *ptr;
                 int n, h, v, x, y, c, j;
+                int block_offset;
                 n = s->nb_blocks[i];
                 c = s->comp_index[i];
                 h = s->h_scount[i];
@@ -807,12 +835,16 @@ static int mjpeg_decode_scan(MJpegDecodeContext *s, int nb_components, int Ah, i
                 x = 0;
                 y = 0;
                 for(j=0;j<n;j++) {
-                    ptr = data[c] +
-                        (((linesize[c] * (v * mb_y + y) * 8) +
-                        (h * mb_x + x) * 8) >> s->avctx->lowres);
+                    block_offset = (((linesize[c] * (v * mb_y + y) * 8) +
+                                     (h * mb_x + x) * 8) >> s->avctx->lowres);
+
                     if(s->interlaced && s->bottom_field)
-                        ptr += linesize[c] >> 1;
+                        block_offset += linesize[c] >> 1;
+                    ptr = data[c] + block_offset;
                     if(!s->progressive) {
+                        if (copy_mb) {
+                            mjpeg_copy_block(ptr, reference_data[c] + block_offset, linesize[c], s->avctx->lowres);
+                        } else {
                         s->dsp.clear_block(s->block);
                         if(decode_block(s, s->block, i,
                                      s->dc_index[i], s->ac_index[i],
@@ -821,6 +853,7 @@ static int mjpeg_decode_scan(MJpegDecodeContext *s, int nb_components, int Ah, i
                             return -1;
                         }
                         s->dsp.idct_put(ptr, linesize[c], s->block);
+                        }
                     } else {
                         int block_idx = s->block_stride[c] * (v * mb_y + y) + (h * mb_x + x);
                         DCTELEM *block = s->blocks[c][block_idx];
@@ -851,29 +884,43 @@ static int mjpeg_decode_scan(MJpegDecodeContext *s, int nb_components, int Ah, i
     return 0;
 }
 
-static int mjpeg_decode_scan_progressive_ac(MJpegDecodeContext *s, int ss, int se, int Ah, int Al){
+static int mjpeg_decode_scan_progressive_ac(MJpegDecodeContext *s, int ss, int se, int Ah, int Al,
+                                            const uint8_t *mb_bitmask, const AVFrame *reference){
     int mb_x, mb_y;
     int EOBRUN = 0;
     int c = s->comp_index[0];
     uint8_t* data = s->picture.data[c];
+    const uint8_t *reference_data = reference ? reference->data[c] : NULL;
     int linesize = s->linesize[c];
     int last_scan = 0;
     int16_t *quant_matrix = s->quant_matrixes[ s->quant_index[c] ];
+    GetBitContext mb_bitmask_gb;
+
+    if (mb_bitmask) {
+        init_get_bits(&mb_bitmask_gb, mb_bitmask, s->mb_width*s->mb_height);
+    }
 
     if(!Al) {
         s->coefs_finished[c] |= (1LL<<(se+1))-(1LL<<ss);
         last_scan = !~s->coefs_finished[c];
     }
 
-    if(s->interlaced && s->bottom_field)
-        data += linesize >> 1;
+    if(s->interlaced && s->bottom_field) {
+        int offset = linesize >> 1;
+        data += offset;
+        reference_data += offset;
+    }
 
     for(mb_y = 0; mb_y < s->mb_height; mb_y++) {
-        uint8_t *ptr = data + (mb_y*linesize*8 >> s->avctx->lowres);
+        int block_offset = (mb_y*linesize*8 >> s->avctx->lowres);
+        uint8_t *ptr = data + block_offset;
         int block_idx = mb_y * s->block_stride[c];
         DCTELEM (*block)[64] = &s->blocks[c][block_idx];
         uint8_t *last_nnz = &s->last_nnz[c][block_idx];
         for(mb_x = 0; mb_x < s->mb_width; mb_x++, block++, last_nnz++) {
+            const int copy_mb = mb_bitmask && !get_bits1(&mb_bitmask_gb);
+
+            if (!copy_mb) {
             int ret;
             if(Ah)
                 ret = decode_block_refinement(s, *block, last_nnz, s->ac_index[0],
@@ -885,16 +932,23 @@ static int mjpeg_decode_scan_progressive_ac(MJpegDecodeContext *s, int ss, int s
                 av_log(s->avctx, AV_LOG_ERROR, "error y=%d x=%d\n", mb_y, mb_x);
                 return -1;
             }
+            }
+
             if(last_scan) {
+                if (copy_mb) {
+                    mjpeg_copy_block(ptr, reference_data + block_offset, linesize, s->avctx->lowres);
+                } else {
                 s->dsp.idct_put(ptr, linesize, *block);
                 ptr += 8 >> s->avctx->lowres;
+                }
             }
         }
     }
     return 0;
 }
 
-int ff_mjpeg_decode_sos(MJpegDecodeContext *s)
+int ff_mjpeg_decode_sos(MJpegDecodeContext *s,
+                        const uint8_t *mb_bitmask, const AVFrame *reference)
 {
     int len, nb_components, i, h, v, predictor, point_transform;
     int index, id;
@@ -996,10 +1050,12 @@ int ff_mjpeg_decode_sos(MJpegDecodeContext *s)
         }
     }else{
         if(s->progressive && predictor) {
-            if(mjpeg_decode_scan_progressive_ac(s, predictor, ilv, prev_shift, point_transform) < 0)
+            if(mjpeg_decode_scan_progressive_ac(s, predictor, ilv, prev_shift, point_transform,
+                                                mb_bitmask, reference) < 0)
                 return -1;
         } else {
-            if(mjpeg_decode_scan(s, nb_components, prev_shift, point_transform) < 0)
+            if(mjpeg_decode_scan(s, nb_components, prev_shift, point_transform,
+                                 mb_bitmask, reference) < 0)
                 return -1;
         }
     }
@@ -1466,7 +1522,7 @@ eoi_parser:
                         av_log(avctx, AV_LOG_WARNING, "Can not process SOS before SOF, skipping\n");
                         break;
                     }
-                    ff_mjpeg_decode_sos(s);
+                    ff_mjpeg_decode_sos(s, NULL, NULL);
                     /* buggy avid puts EOI every 10-20th frame */
                     /* if restart period is over process EOI */
                     if ((s->buggy_avid && !s->interlaced) || s->restart_interval)
diff --git a/libavcodec/mjpegdec.h b/libavcodec/mjpegdec.h
index a190fdadf8..63bcb93046 100644
--- a/libavcodec/mjpegdec.h
+++ b/libavcodec/mjpegdec.h
@@ -115,6 +115,7 @@ int ff_mjpeg_decode_frame(AVCodecContext *avctx,
 int ff_mjpeg_decode_dqt(MJpegDecodeContext *s);
 int ff_mjpeg_decode_dht(MJpegDecodeContext *s);
 int ff_mjpeg_decode_sof(MJpegDecodeContext *s);
-int ff_mjpeg_decode_sos(MJpegDecodeContext *s);
+int ff_mjpeg_decode_sos(MJpegDecodeContext *s,
+                        const uint8_t *mb_bitmask, const AVFrame *reference);
 
 #endif /* AVCODEC_MJPEGDEC_H */