snow: yuv444 support

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

libavcodec/snow.c

@@ -328,7 +328,7 @@ void ff_snow_pred_block(SnowContext *s, uint8_t *dst, uint8_t *tmp, int stride,
         }
     }else{
         uint8_t *src= s->last_picture[block->ref].data[plane_index];
-        const int scale= plane_index ?  s->mv_scale : 2*s->mv_scale;
+        const int scale= plane_index ?  (2*s->mv_scale)>>s->chroma_h_shift : 2*s->mv_scale;
         int mx= block->mx*scale;
         int my= block->my*scale;
         const int dx= mx&15;
@@ -342,6 +342,9 @@ void ff_snow_pred_block(SnowContext *s, uint8_t *dst, uint8_t *tmp, int stride,
             s->dsp.emulated_edge_mc(tmp + MB_SIZE, src, stride, b_w+HTAPS_MAX-1, b_h+HTAPS_MAX-1, sx, sy, w, h);
             src= tmp + MB_SIZE;
         }
+
+        av_assert2(s->chroma_h_shift == s->chroma_v_shift); // only one mv_scale
+
 //        assert(b_w == b_h || 2*b_w == b_h || b_w == 2*b_h);
 //        assert(!(b_w&(b_w-1)));
         assert(b_w>1 && b_h>1);
@@ -513,8 +516,8 @@ static void halfpel_interpol(SnowContext *s, uint8_t *halfpel[4][4], AVFrame *fr
 
     for(p=0; p<3; p++){
         int is_chroma= !!p;
-        int w= s->avctx->width  >>is_chroma;
-        int h= s->avctx->height >>is_chroma;
+        int w= is_chroma ? s->avctx->width >>s->chroma_h_shift : s->avctx->width;
+        int h= is_chroma ? s->avctx->height>>s->chroma_v_shift : s->avctx->height;
         int ls= frame->linesize[p];
         uint8_t *src= frame->data[p];
 
@@ -573,11 +576,11 @@ int ff_snow_frame_start(SnowContext *s){
                           s->current_picture.linesize[0], w   , h   ,
                           EDGE_WIDTH  , EDGE_WIDTH  , EDGE_TOP | EDGE_BOTTOM);
         s->dsp.draw_edges(s->current_picture.data[1],
-                          s->current_picture.linesize[1], w>>1, h>>1,
-                          EDGE_WIDTH/2, EDGE_WIDTH/2, EDGE_TOP | EDGE_BOTTOM);
+                          s->current_picture.linesize[1], w>>s->chroma_h_shift, h>>s->chroma_v_shift,
+                          EDGE_WIDTH>>s->chroma_h_shift, EDGE_WIDTH>>s->chroma_v_shift, EDGE_TOP | EDGE_BOTTOM);
         s->dsp.draw_edges(s->current_picture.data[2],
-                          s->current_picture.linesize[2], w>>1, h>>1,
-                          EDGE_WIDTH/2, EDGE_WIDTH/2, EDGE_TOP | EDGE_BOTTOM);
+                          s->current_picture.linesize[2], w>>s->chroma_h_shift, h>>s->chroma_v_shift,
+                          EDGE_WIDTH>>s->chroma_h_shift, EDGE_WIDTH>>s->chroma_v_shift, EDGE_TOP | EDGE_BOTTOM);
     }
 
     ff_snow_release_buffer(s->avctx);

libavcodec/snow.h

@@ -404,20 +404,21 @@ static av_always_inline void predict_slice(SnowContext *s, IDWTELEM *buf, int pl
     const int mb_h= s->b_height << s->block_max_depth;
     int x, y, mb_x;
     int block_size = MB_SIZE >> s->block_max_depth;
-    int block_w    = plane_index ? block_size/2 : block_size;
-    const uint8_t *obmc  = plane_index ? ff_obmc_tab[s->block_max_depth+1] : ff_obmc_tab[s->block_max_depth];
-    const int obmc_stride= plane_index ? block_size : 2*block_size;
+    int block_w    = plane_index ? block_size>>s->chroma_h_shift : block_size;
+    int block_h    = plane_index ? block_size>>s->chroma_v_shift : block_size;
+    const uint8_t *obmc  = plane_index ? ff_obmc_tab[s->block_max_depth+s->chroma_h_shift] : ff_obmc_tab[s->block_max_depth];
+    const int obmc_stride= plane_index ? (2*block_size)>>s->chroma_h_shift : 2*block_size;
     int ref_stride= s->current_picture.linesize[plane_index];
     uint8_t *dst8= s->current_picture.data[plane_index];
     int w= p->width;
     int h= p->height;
-
+    av_assert2(s->chroma_h_shift == s->chroma_v_shift); // obmc params assume squares
     if(s->keyframe || (s->avctx->debug&512)){
         if(mb_y==mb_h)
             return;
 
         if(add){
-            for(y=block_w*mb_y; y<FFMIN(h,block_w*(mb_y+1)); y++){
+            for(y=block_h*mb_y; y<FFMIN(h,block_h*(mb_y+1)); y++){
                 for(x=0; x<w; x++){
                     int v= buf[x + y*w] + (128<<FRAC_BITS) + (1<<(FRAC_BITS-1));
                     v >>= FRAC_BITS;
@@ -426,7 +427,7 @@ static av_always_inline void predict_slice(SnowContext *s, IDWTELEM *buf, int pl
                 }
             }
         }else{
-            for(y=block_w*mb_y; y<FFMIN(h,block_w*(mb_y+1)); y++){
+            for(y=block_h*mb_y; y<FFMIN(h,block_h*(mb_y+1)); y++){
                 for(x=0; x<w; x++){
                     buf[x + y*w]-= 128<<FRAC_BITS;
                 }
@@ -439,8 +440,8 @@ static av_always_inline void predict_slice(SnowContext *s, IDWTELEM *buf, int pl
     for(mb_x=0; mb_x<=mb_w; mb_x++){
         add_yblock(s, 0, NULL, buf, dst8, obmc,
                    block_w*mb_x - block_w/2,
-                   block_w*mb_y - block_w/2,
-                   block_w, block_w,
+                   block_h*mb_y - block_h/2,
+                   block_w, block_h,
                    w, h,
                    w, ref_stride, obmc_stride,
                    mb_x - 1, mb_y - 1,
@@ -460,6 +461,7 @@ static inline void set_blocks(SnowContext *s, int level, int x, int y, int l, in
     const int rem_depth= s->block_max_depth - level;
     const int index= (x + y*w) << rem_depth;
     const int block_w= 1<<rem_depth;
+    const int block_h= 1<<rem_depth; //FIXME "w!=h"
     BlockNode block;
     int i,j;
 
@@ -472,7 +474,7 @@ static inline void set_blocks(SnowContext *s, int level, int x, int y, int l, in
     block.type= type;
     block.level= level;
 
-    for(j=0; j<block_w; j++){
+    for(j=0; j<block_h; j++){
         for(i=0; i<block_w; i++){
             s->block[index + i + j*w]= block;
         }
@@ -480,10 +482,11 @@ static inline void set_blocks(SnowContext *s, int level, int x, int y, int l, in
 }
 
 static inline void init_ref(MotionEstContext *c, uint8_t *src[3], uint8_t *ref[3], uint8_t *ref2[3], int x, int y, int ref_index){
+    SnowContext *s = c->avctx->priv_data;
     const int offset[3]= {
           y*c->  stride + x,
-        ((y*c->uvstride + x)>>1),
-        ((y*c->uvstride + x)>>1),
+        ((y*c->uvstride + x)>>s->chroma_h_shift),
+        ((y*c->uvstride + x)>>s->chroma_h_shift),
     };
     int i;
     for(i=0; i<3; i++){

libavcodec/snowdec.c

@@ -41,9 +41,10 @@ static av_always_inline void predict_slice_buffered(SnowContext *s, slice_buffer
     const int mb_h= s->b_height << s->block_max_depth;
     int x, y, mb_x;
     int block_size = MB_SIZE >> s->block_max_depth;
-    int block_w    = plane_index ? block_size/2 : block_size;
-    const uint8_t *obmc  = plane_index ? ff_obmc_tab[s->block_max_depth+1] : ff_obmc_tab[s->block_max_depth];
-    int obmc_stride= plane_index ? block_size : 2*block_size;
+    int block_w    = plane_index ? block_size>>s->chroma_h_shift : block_size;
+    int block_h    = plane_index ? block_size>>s->chroma_v_shift : block_size;
+    const uint8_t *obmc  = plane_index ? ff_obmc_tab[s->block_max_depth+s->chroma_h_shift] : ff_obmc_tab[s->block_max_depth];
+    int obmc_stride= plane_index ? (2*block_size)>>s->chroma_h_shift : 2*block_size;
     int ref_stride= s->current_picture.linesize[plane_index];
     uint8_t *dst8= s->current_picture.data[plane_index];
     int w= p->width;
@@ -54,7 +55,7 @@ static av_always_inline void predict_slice_buffered(SnowContext *s, slice_buffer
             return;
 
         if(add){
-            for(y=block_w*mb_y; y<FFMIN(h,block_w*(mb_y+1)); y++){
+            for(y=block_h*mb_y; y<FFMIN(h,block_h*(mb_y+1)); y++){
 //                DWTELEM * line = slice_buffer_get_line(sb, y);
                 IDWTELEM * line = sb->line[y];
                 for(x=0; x<w; x++){
@@ -66,7 +67,7 @@ static av_always_inline void predict_slice_buffered(SnowContext *s, slice_buffer
                 }
             }
         }else{
-            for(y=block_w*mb_y; y<FFMIN(h,block_w*(mb_y+1)); y++){
+            for(y=block_h*mb_y; y<FFMIN(h,block_h*(mb_y+1)); y++){
 //                DWTELEM * line = slice_buffer_get_line(sb, y);
                 IDWTELEM * line = sb->line[y];
                 for(x=0; x<w; x++){
@@ -82,8 +83,8 @@ static av_always_inline void predict_slice_buffered(SnowContext *s, slice_buffer
     for(mb_x=0; mb_x<=mb_w; mb_x++){
         add_yblock(s, 1, sb, old_buffer, dst8, obmc,
                    block_w*mb_x - block_w/2,
-                   block_w*mb_y - block_w/2,
-                   block_w, block_w,
+                   block_h*mb_y - block_h/2,
+                   block_w, block_h,
                    w, h,
                    w, ref_stride, obmc_stride,
                    mb_x - 1, mb_y - 1,
@@ -289,6 +290,18 @@ static int decode_header(SnowContext *s){
         s->colorspace_type= get_symbol(&s->c, s->header_state, 0);
         s->chroma_h_shift= get_symbol(&s->c, s->header_state, 0);
         s->chroma_v_shift= get_symbol(&s->c, s->header_state, 0);
+
+        if(s->chroma_h_shift == 1 && s->chroma_v_shift==1){
+            s->avctx->pix_fmt= PIX_FMT_YUV420P;
+        }else if(s->chroma_h_shift == 0 && s->chroma_v_shift==0){
+            s->avctx->pix_fmt= PIX_FMT_YUV444P;
+        } else {
+            av_log(s, AV_LOG_ERROR, "unsupported color subsample mode %d %d\n", s->chroma_h_shift, s->chroma_v_shift);
+            s->chroma_h_shift = s->chroma_v_shift = 1;
+            s->avctx->pix_fmt= PIX_FMT_YUV420P;
+            return AVERROR_INVALIDDATA;
+        }
+
         s->spatial_scalability= get_rac(&s->c, s->header_state);
 //        s->rate_scalability= get_rac(&s->c, s->header_state);
         GET_S(s->max_ref_frames, tmp < (unsigned)MAX_REF_FRAMES)
@@ -334,10 +347,6 @@ static int decode_header(SnowContext *s){
         return -1;
     }
 
-    if (s->chroma_h_shift != 1 || s->chroma_v_shift != 1) {
-        av_log(s->avctx, AV_LOG_ERROR, "Invalid chroma shift\n");
-        return AVERROR_PATCHWELCOME;
-    }
 
     s->qlog           += get_symbol(&s->c, s->header_state, 1);
     s->mv_scale       += get_symbol(&s->c, s->header_state, 1);
@@ -354,8 +363,6 @@ static int decode_header(SnowContext *s){
 
 static av_cold int decode_init(AVCodecContext *avctx)
 {
-    avctx->pix_fmt= PIX_FMT_YUV420P;
-
     ff_snow_common_init(avctx);
 
     return 0;
@@ -452,7 +459,8 @@ static int decode_frame(AVCodecContext *avctx, void *data, int *data_size, AVPac
         {
         const int mb_h= s->b_height << s->block_max_depth;
         const int block_size = MB_SIZE >> s->block_max_depth;
-        const int block_w    = plane_index ? block_size/2 : block_size;
+        const int block_w    = plane_index ? block_size>>s->chroma_h_shift : block_size;
+        const int block_h    = plane_index ? block_size>>s->chroma_v_shift : block_size;
         int mb_y;
         DWTCompose cs[MAX_DECOMPOSITIONS];
         int yd=0, yq=0;
@@ -462,11 +470,12 @@ static int decode_frame(AVCodecContext *avctx, void *data, int *data_size, AVPac
         ff_spatial_idwt_buffered_init(cs, &s->sb, w, h, 1, s->spatial_decomposition_type, s->spatial_decomposition_count);
         for(mb_y=0; mb_y<=mb_h; mb_y++){
 
-            int slice_starty = block_w*mb_y;
-            int slice_h = block_w*(mb_y+1);
+            int slice_starty = block_h*mb_y;
+            int slice_h = block_h*(mb_y+1);
+
             if (!(s->keyframe || s->avctx->debug&512)){
-                slice_starty = FFMAX(0, slice_starty - (block_w >> 1));
-                slice_h -= (block_w >> 1);
+                slice_starty = FFMAX(0, slice_starty - (block_h >> 1));
+                slice_h -= (block_h >> 1);
             }
 
             for(level=0; level<s->spatial_decomposition_count; level++){
@@ -477,11 +486,11 @@ static int decode_frame(AVCodecContext *avctx, void *data, int *data_size, AVPac
                     int our_mb_start = mb_y;
                     int our_mb_end = (mb_y + 1);
                     const int extra= 3;
-                    start_y = (mb_y ? ((block_w * our_mb_start) >> (s->spatial_decomposition_count - level)) + s->spatial_decomposition_count - level + extra: 0);
-                    end_y = (((block_w * our_mb_end) >> (s->spatial_decomposition_count - level)) + s->spatial_decomposition_count - level + extra);
+                    start_y = (mb_y ? ((block_h * our_mb_start) >> (s->spatial_decomposition_count - level)) + s->spatial_decomposition_count - level + extra: 0);
+                    end_y = (((block_h * our_mb_end) >> (s->spatial_decomposition_count - level)) + s->spatial_decomposition_count - level + extra);
                     if (!(s->keyframe || s->avctx->debug&512)){
-                        start_y = FFMAX(0, start_y - (block_w >> (1+s->spatial_decomposition_count - level)));
-                        end_y = FFMAX(0, end_y - (block_w >> (1+s->spatial_decomposition_count - level)));
+                        start_y = FFMAX(0, start_y - (block_h >> (1+s->spatial_decomposition_count - level)));
+                        end_y = FFMAX(0, end_y - (block_h >> (1+s->spatial_decomposition_count - level)));
                     }
                     start_y = FFMIN(b->height, start_y);
                     end_y = FFMIN(b->height, end_y);