avcodec/snowdec: AVMotionVector support

Signed-off-by: Michael Niedermayer <michaelni@gmx.at>

1 files changed, 36 insertions, 0 deletions

diff --git a/libavcodec/snowdec.c b/libavcodec/snowdec.c
index 462ecf1953..58f9ed894b 100644
--- a/libavcodec/snowdec.c
+++ b/libavcodec/snowdec.c
@@ -87,6 +87,26 @@ static av_always_inline void predict_slice_buffered(SnowContext *s, slice_buffer
                    mb_x - 1, mb_y - 1,
                    add, 0, plane_index);
     }
+
+    if(s->avmv && mb_y < mb_h && plane_index == 0)
+        for(mb_x=0; mb_x<mb_w; mb_x++){
+            AVMotionVector *avmv = s->avmv + (s->avmv_index++);
+            const int b_width = s->b_width  << s->block_max_depth;
+            const int b_stride= b_width;
+            BlockNode *bn= &s->block[mb_x + mb_y*b_stride];
+
+            if (bn->type)
+                continue;
+
+            avmv->w = block_w;
+            avmv->h = block_h;
+            avmv->dst_x = block_w*mb_x - block_w/2;
+            avmv->dst_y = block_h*mb_y - block_h/2;
+            avmv->src_x = avmv->dst_x + (bn->mx * s->mv_scale)/8;
+            avmv->src_y = avmv->dst_y + (bn->my * s->mv_scale)/8;
+            avmv->source= -1 - bn->ref;
+            avmv->flags = 0;
+        }
 }
 
 static inline void decode_subband_slice_buffered(SnowContext *s, SubBand *b, slice_buffer * sb, int start_y, int h, int save_state[1]){
@@ -455,6 +475,12 @@ static int decode_frame(AVCodecContext *avctx, void *data, int *got_frame,
                s->spatial_decomposition_count
               );
 
+    av_assert0(!s->avmv);
+    if (s->avctx->flags2 & CODEC_FLAG2_EXPORT_MVS) {
+        s->avmv = av_malloc_array(s->b_width * s->b_height, sizeof(AVMotionVector) << (s->block_max_depth*2));
+    }
+    s->avmv_index = 0;
+
     if ((res = decode_blocks(s)) < 0)
         return res;
 
@@ -573,6 +599,16 @@ static int decode_frame(AVCodecContext *avctx, void *data, int *got_frame,
         res = av_frame_ref(picture, s->current_picture);
     else
         res = av_frame_ref(picture, s->mconly_picture);
+    if (res >= 0 && s->avmv_index) {
+        AVFrameSideData *sd;
+
+        sd = av_frame_new_side_data(picture, AV_FRAME_DATA_MOTION_VECTORS, s->avmv_index * sizeof(AVMotionVector));
+        if (!sd)
+            return AVERROR(ENOMEM);
+        memcpy(sd->data, s->avmv, s->avmv_index * sizeof(AVMotionVector));
+    }
+
+    av_freep(&s->avmv);
 
     if (res < 0)
         return res;