motion_est.c

arm平台下的H264编码和解码源代码

/*
 * Motion estimation 
 * Copyright (c) 2000,2001 Fabrice Bellard.
 * Copyright (c) 2002-2004 Michael Niedermayer
 * 
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 * new Motion Estimation (X1/EPZS) by Michael Niedermayer <[e-mail]>
 */
 
/**
 * @file motion_est.c
 * Motion estimation.
 */
 
#include <stdlib.h>
#include <stdio.h>
#include <limits.h>
#include "avcodec.h"
#include "dsputil.h"
#include "mpegvideo.h"

//#undef NDEBUG
//#include <assert.h>

#define SQ(a) ((a)*(a))

#define P_LEFT P[1]
#define P_TOP P[2]
#define P_TOPRIGHT P[3]
#define P_MEDIAN P[4]
#define P_MV1 P[9]

static inline int sad_hpel_motion_search(MpegEncContext * s,
				  int *mx_ptr, int *my_ptr, int dmin,
                                  int pred_x, int pred_y, uint8_t *src_data[3],
                                  uint8_t *ref_data[6], int stride, int uvstride,
                                  int size, int h, uint8_t * const mv_penalty);

static inline int update_map_generation(MpegEncContext * s)
{
    s->me.map_generation+= 1<<(ME_MAP_MV_BITS*2);
    if(s->me.map_generation==0){
        s->me.map_generation= 1<<(ME_MAP_MV_BITS*2);
        memset(s->me.map, 0, sizeof(uint32_t)*ME_MAP_SIZE);
    }
    return s->me.map_generation;
}

/* shape adaptive search stuff */
typedef struct Minima{
    int height;
    int x, y;
    int checked;
}Minima;

static int minima_cmp(const void *a, const void *b){
    const Minima *da = (const Minima *) a;
    const Minima *db = (const Minima *) b;
    
    return da->height - db->height;
}
                                  
/* SIMPLE */
#define RENAME(a) simple_ ## a

#define CMP(d, x, y, size)\
d = cmp(s, src_y, (ref_y) + (x) + (y)*(stride), stride, h);

#define CMP_HPEL(d, dx, dy, x, y, size)\
{\
    const int dxy= (dx) + 2*(dy);\
    hpel_put[0][dxy](s->me.scratchpad, (ref_y) + (x) + (y)*(stride), stride, h);\
    d = cmp_sub(s, s->me.scratchpad, src_y, stride, h);\
}


#define CMP_QPEL(d, dx, dy, x, y, size)\
{\
    const int dxy= (dx) + 4*(dy);\
    qpel_put[0][dxy](s->me.scratchpad, (ref_y) + (x) + (y)*(stride), stride);\
    d = cmp_sub(s, s->me.scratchpad, src_y, stride, h);\
}

#include "motion_est_template.c"
#undef RENAME
#undef CMP
#undef CMP_HPEL
#undef CMP_QPEL
#undef INIT

/* SIMPLE CHROMA */
#define RENAME(a) simple_chroma_ ## a

#define CMP(d, x, y, size)\
d = cmp(s, src_y, (ref_y) + (x) + (y)*(stride), stride, h);\
if(chroma_cmp){\
    int dxy= ((x)&1) + 2*((y)&1);\
    int c= ((x)>>1) + ((y)>>1)*uvstride;\
\
    chroma_hpel_put[0][dxy](s->me.scratchpad, ref_u + c, uvstride, h>>1);\
    d += chroma_cmp(s, s->me.scratchpad, src_u, uvstride, h>>1);\
    chroma_hpel_put[0][dxy](s->me.scratchpad, ref_v + c, uvstride, h>>1);\
    d += chroma_cmp(s, s->me.scratchpad, src_v, uvstride, h>>1);\
}

#define CMP_HPEL(d, dx, dy, x, y, size)\
{\
    const int dxy= (dx) + 2*(dy);\
    hpel_put[0][dxy](s->me.scratchpad, (ref_y) + (x) + (y)*(stride), stride, h);\
    d = cmp_sub(s, s->me.scratchpad, src_y, stride, h);\
    if(chroma_cmp_sub){\
        int cxy= (dxy) | ((x)&1) | (2*((y)&1));\
        int c= ((x)>>1) + ((y)>>1)*uvstride;\
        chroma_hpel_put[0][cxy](s->me.scratchpad, ref_u + c, uvstride, h>>1);\
        d += chroma_cmp_sub(s, s->me.scratchpad, src_u, uvstride, h>>1);\
        chroma_hpel_put[0][cxy](s->me.scratchpad, ref_v + c, uvstride, h>>1);\
        d += chroma_cmp_sub(s, s->me.scratchpad, src_v, uvstride, h>>1);\
    }\
}

#define CMP_QPEL(d, dx, dy, x, y, size)\
{\
    const int dxy= (dx) + 4*(dy);\
    qpel_put[0][dxy](s->me.scratchpad, (ref_y) + (x) + (y)*(stride), stride);\
    d = cmp_sub(s, s->me.scratchpad, src_y, stride, h);\
    if(chroma_cmp_sub){\
        int cxy, c;\
        int cx= (4*(x) + (dx))/2;\
        int cy= (4*(y) + (dy))/2;\
        cx= (cx>>1)|(cx&1);\
        cy= (cy>>1)|(cy&1);\
        cxy= (cx&1) + 2*(cy&1);\
        c= ((cx)>>1) + ((cy)>>1)*uvstride;\
        chroma_hpel_put[0][cxy](s->me.scratchpad, ref_u + c, uvstride, h>>1);\
        d += chroma_cmp_sub(s, s->me.scratchpad, src_u, uvstride, h>>1);\
        chroma_hpel_put[0][cxy](s->me.scratchpad, ref_v + c, uvstride, h>>1);\
        d += chroma_cmp_sub(s, s->me.scratchpad, src_v, uvstride, h>>1);\
    }\
}

#include "motion_est_template.c"
#undef RENAME
#undef CMP
#undef CMP_HPEL
#undef CMP_QPEL
#undef INIT

/* SIMPLE DIRECT HPEL */
#define RENAME(a) simple_direct_hpel_ ## a
//FIXME precalc divisions stuff

#define CMP_DIRECT(d, dx, dy, x, y, size, cmp_func)\
if((x) >= xmin && 2*(x) + (dx) <= 2*xmax && (y) >= ymin && 2*(y) + (dy) <= 2*ymax){\
    const int hx= 2*(x) + (dx);\
    const int hy= 2*(y) + (dy);\
    if(s->mv_type==MV_TYPE_8X8){\
        int i;\
        for(i=0; i<4; i++){\
            int fx = s->me.direct_basis_mv[i][0] + hx;\
            int fy = s->me.direct_basis_mv[i][1] + hy;\
            int bx = hx ? fx - s->me.co_located_mv[i][0] : s->me.co_located_mv[i][0]*(time_pb - time_pp)/time_pp + (i &1)*16;\
            int by = hy ? fy - s->me.co_located_mv[i][1] : s->me.co_located_mv[i][1]*(time_pb - time_pp)/time_pp + (i>>1)*16;\
            int fxy= (fx&1) + 2*(fy&1);\
            int bxy= (bx&1) + 2*(by&1);\
\
            uint8_t *dst= s->me.scratchpad + 8*(i&1) + 8*stride*(i>>1);\
            hpel_put[1][fxy](dst, (ref_y ) + (fx>>1) + (fy>>1)*(stride), stride, 8);\
            hpel_avg[1][bxy](dst, (ref_data[3]) + (bx>>1) + (by>>1)*(stride), stride, 8);\
        }\
    }else{\
        int fx = s->me.direct_basis_mv[0][0] + hx;\
        int fy = s->me.direct_basis_mv[0][1] + hy;\
        int bx = hx ? fx - s->me.co_located_mv[0][0] : (s->me.co_located_mv[0][0]*(time_pb - time_pp)/time_pp);\
        int by = hy ? fy - s->me.co_located_mv[0][1] : (s->me.co_located_mv[0][1]*(time_pb - time_pp)/time_pp);\
        int fxy= (fx&1) + 2*(fy&1);\
        int bxy= (bx&1) + 2*(by&1);\
        \
        assert((fx>>1) + 16*s->mb_x >= -16);\
        assert((fy>>1) + 16*s->mb_y >= -16);\
        assert((fx>>1) + 16*s->mb_x <= s->width);\
        assert((fy>>1) + 16*s->mb_y <= s->height);\
        assert((bx>>1) + 16*s->mb_x >= -16);\
        assert((by>>1) + 16*s->mb_y >= -16);\
        assert((bx>>1) + 16*s->mb_x <= s->width);\
        assert((by>>1) + 16*s->mb_y <= s->height);\
\
        hpel_put[0][fxy](s->me.scratchpad, (ref_y ) + (fx>>1) + (fy>>1)*(stride), stride, 16);\
        hpel_avg[0][bxy](s->me.scratchpad, (ref_data[3]) + (bx>>1) + (by>>1)*(stride), stride, 16);\
    }\
    d = cmp_func(s, s->me.scratchpad, src_y, stride, 16);\
}else\
    d= 256*256*256*32;


#define CMP_HPEL(d, dx, dy, x, y, size)\
    CMP_DIRECT(d, dx, dy, x, y, size, cmp_sub)

#define CMP(d, x, y, size)\
    CMP_DIRECT(d, 0, 0, x, y, size, cmp)
    
#include "motion_est_template.c"
#undef RENAME
#undef CMP
#undef CMP_HPEL
#undef CMP_QPEL
#undef INIT
#undef CMP_DIRECT

/* SIMPLE DIRECT QPEL */
#define RENAME(a) simple_direct_qpel_ ## a

#define CMP_DIRECT(d, dx, dy, x, y, size, cmp_func)\
if((x) >= xmin && 4*(x) + (dx) <= 4*xmax && (y) >= ymin && 4*(y) + (dy) <= 4*ymax){\
    const int qx= 4*(x) + (dx);\
    const int qy= 4*(y) + (dy);\
    if(s->mv_type==MV_TYPE_8X8){\
        int i;\
        for(i=0; i<4; i++){\
            int fx = s->me.direct_basis_mv[i][0] + qx;\
            int fy = s->me.direct_basis_mv[i][1] + qy;\
            int bx = qx ? fx - s->me.co_located_mv[i][0] : s->me.co_located_mv[i][0]*(time_pb - time_pp)/time_pp + (i &1)*16;\
            int by = qy ? fy - s->me.co_located_mv[i][1] : s->me.co_located_mv[i][1]*(time_pb - time_pp)/time_pp + (i>>1)*16;\
            int fxy= (fx&3) + 4*(fy&3);\
            int bxy= (bx&3) + 4*(by&3);\
\
            uint8_t *dst= s->me.scratchpad + 8*(i&1) + 8*stride*(i>>1);\
            qpel_put[1][fxy](dst, (ref_y ) + (fx>>2) + (fy>>2)*(stride), stride);\
            qpel_avg[1][bxy](dst, (ref_data[3]) + (bx>>2) + (by>>2)*(stride), stride);\
        }\
    }else{\
        int fx = s->me.direct_basis_mv[0][0] + qx;\
        int fy = s->me.direct_basis_mv[0][1] + qy;\
        int bx = qx ? fx - s->me.co_located_mv[0][0] : s->me.co_located_mv[0][0]*(time_pb - time_pp)/time_pp;\
        int by = qy ? fy - s->me.co_located_mv[0][1] : s->me.co_located_mv[0][1]*(time_pb - time_pp)/time_pp;\
        int fxy= (fx&3) + 4*(fy&3);\
        int bxy= (bx&3) + 4*(by&3);\
\
        qpel_put[1][fxy](s->me.scratchpad               , (ref_y ) + (fx>>2) + (fy>>2)*(stride)               , stride);\
        qpel_put[1][fxy](s->me.scratchpad + 8           , (ref_y ) + (fx>>2) + (fy>>2)*(stride) + 8           , stride);\
        qpel_put[1][fxy](s->me.scratchpad     + 8*stride, (ref_y ) + (fx>>2) + (fy>>2)*(stride)     + 8*stride, stride);\
        qpel_put[1][fxy](s->me.scratchpad + 8 + 8*stride, (ref_y ) + (fx>>2) + (fy>>2)*(stride) + 8 + 8*stride, stride);\
        qpel_avg[1][bxy](s->me.scratchpad               , (ref_data[3]) + (bx>>2) + (by>>2)*(stride)               , stride);\
        qpel_avg[1][bxy](s->me.scratchpad + 8           , (ref_data[3]) + (bx>>2) + (by>>2)*(stride) + 8           , stride);\
        qpel_avg[1][bxy](s->me.scratchpad     + 8*stride, (ref_data[3]) + (bx>>2) + (by>>2)*(stride)     + 8*stride, stride);\
        qpel_avg[1][bxy](s->me.scratchpad + 8 + 8*stride, (ref_data[3]) + (bx>>2) + (by>>2)*(stride) + 8 + 8*stride, stride);\
    }\
    d = cmp_func(s, s->me.scratchpad, src_y, stride, 16);\
}else\
    d= 256*256*256*32;


#define CMP_QPEL(d, dx, dy, x, y, size)\
    CMP_DIRECT(d, dx, dy, x, y, size, cmp_sub)

#define CMP(d, x, y, size)\
    CMP_DIRECT(d, 0, 0, x, y, size, cmp)

#include "motion_est_template.c"
#undef RENAME
#undef CMP
#undef CMP_HPEL
#undef CMP_QPEL
#undef INIT
#undef CMP__DIRECT

static inline int get_penalty_factor(MpegEncContext *s, int type){
    switch(type&0xFF){
    default:
    case FF_CMP_SAD:
        return s->qscale*2;
    case FF_CMP_DCT:
        return s->qscale*3;
    case FF_CMP_SATD:
        return s->qscale*6;
    case FF_CMP_SSE:
        return s->qscale*s->qscale*2;
    case FF_CMP_BIT:
        return 1;
    case FF_CMP_RD:
    case FF_CMP_PSNR:
        return (s->qscale*s->qscale*185 + 64)>>7;
    }
}

void ff_init_me(MpegEncContext *s){
    ff_set_cmp(&s->dsp, s->dsp.me_pre_cmp, s->avctx->me_pre_cmp);
    ff_set_cmp(&s->dsp, s->dsp.me_cmp, s->avctx->me_cmp);
    ff_set_cmp(&s->dsp, s->dsp.me_sub_cmp, s->avctx->me_sub_cmp);
    ff_set_cmp(&s->dsp, s->dsp.mb_cmp, s->avctx->mb_cmp);

    if(s->flags&CODEC_FLAG_QPEL){
        if(s->avctx->me_sub_cmp&FF_CMP_CHROMA)
            s->me.sub_motion_search= simple_chroma_qpel_motion_search;
        else
            s->me.sub_motion_search= simple_qpel_motion_search;
    }else{
        if(s->avctx->me_sub_cmp&FF_CMP_CHROMA)
            s->me.sub_motion_search= simple_chroma_hpel_motion_search;
        else if(   s->avctx->me_sub_cmp == FF_CMP_SAD 
                && s->avctx->    me_cmp == FF_CMP_SAD 
                && s->avctx->    mb_cmp == FF_CMP_SAD)
            s->me.sub_motion_search= sad_hpel_motion_search; // 2050 vs. 2450 cycles
        else
            s->me.sub_motion_search= simple_hpel_motion_search;
    }

    if(s->avctx->me_cmp&FF_CMP_CHROMA){
        s->me.motion_search[0]= simple_chroma_epzs_motion_search;
        s->me.motion_search[1]= simple_chroma_epzs_motion_search4;
        s->me.motion_search[4]= simple_chroma_epzs_motion_search2;
    }else{
        s->me.motion_search[0]= simple_epzs_motion_search;
        s->me.motion_search[1]= simple_epzs_motion_search4;
        s->me.motion_search[4]= simple_epzs_motion_search2;
    }
    
    if(s->avctx->me_pre_cmp&FF_CMP_CHROMA){
        s->me.pre_motion_search= simple_chroma_epzs_motion_search;
    }else{
        s->me.pre_motion_search= simple_epzs_motion_search;
    }
    
    if(s->flags&CODEC_FLAG_QPEL){
        if(s->avctx->mb_cmp&FF_CMP_CHROMA)
            s->me.get_mb_score= simple_chroma_qpel_get_mb_score;
        else
            s->me.get_mb_score= simple_qpel_get_mb_score;
    }else{
        if(s->avctx->mb_cmp&FF_CMP_CHROMA)
            s->me.get_mb_score= simple_chroma_hpel_get_mb_score;
        else
            s->me.get_mb_score= simple_hpel_get_mb_score;
    }
}
      
#if 0
static int pix_dev(uint8_t * pix, int line_size, int mean)
{
    int s, i, j;

    s = 0;
    for (i = 0; i < 16; i++) {
	for (j = 0; j < 16; j += 8) {
	    s += ABS(pix[0]-mean);
	    s += ABS(pix[1]-mean);
	    s += ABS(pix[2]-mean);
	    s += ABS(pix[3]-mean);
	    s += ABS(pix[4]-mean);
	    s += ABS(pix[5]-mean);
	    s += ABS(pix[6]-mean);
	    s += ABS(pix[7]-mean);
	    pix += 8;
	}
	pix += line_size - 16;