t264enc.c

T264是中国的视频编码自由组织合力开发的264编解码程序

    if (para->flags & USE_QUARTPEL)
        para->flags |= USE_HALFPEL;

    t = T264_malloc(sizeof(T264_t), CACHE_SIZE);
    memset(t, 0, sizeof(T264_t));

    t->mb_width  = para->width >> 4;
    t->mb_height = para->height >> 4;
    t->mb_stride = t->mb_width;
    t->width  = t->mb_width << 4;
    t->height = t->mb_height << 4;
    t->edged_width = t->width + 2 * EDGED_WIDTH;
    t->edged_height = t->height + 2 * EDGED_HEIGHT;
    t->qp_y   = para->qp;
    t->flags  = para->flags;

    t->stride    = t->width;
    t->stride_uv = t->width >> 1;
    t->edged_stride = t->edged_width;
    t->edged_stride_uv = t->edged_width >> 1;

    t->bs = T264_malloc(sizeof(bs_t), CACHE_SIZE);
    t->bs_buf = T264_malloc(t->width * t->height << 1, CACHE_SIZE);
    para->direct_flag = 1;  /* force direct mode */
    if (para->b_num)
        para->ref_num ++;

    for(i = 0 ; i < para->ref_num + 1 ; i ++)
    {
        uint8_t* p = T264_malloc(t->edged_width * t->edged_height + (t->edged_width * t->edged_height >> 1), CACHE_SIZE);
        t->refn[i].Y[0] = p + EDGED_HEIGHT * t->edged_width + EDGED_WIDTH;
        t->refn[i].U = p + t->edged_width * t->edged_height + (t->edged_width * EDGED_HEIGHT >> 2) + (EDGED_WIDTH >> 1);
        t->refn[i].V = p + t->edged_width * t->edged_height + (t->edged_width * t->edged_height >> 2) + (t->edged_width * EDGED_HEIGHT >> 2) + (EDGED_WIDTH >> 1);
        t->refn[i].mb = T264_malloc(t->mb_height * t->mb_width * sizeof(T264_mb_context_t), CACHE_SIZE);
        p = T264_malloc(t->edged_width * t->edged_height * 3, CACHE_SIZE);
        t->refn[i].Y[1] = p + EDGED_HEIGHT * t->edged_width + EDGED_WIDTH;
        t->refn[i].Y[2] = t->refn[i].Y[1] + t->edged_width * t->edged_height;
        t->refn[i].Y[3] = t->refn[i].Y[2] + t->edged_width * t->edged_height;
    }
    for(i = 0 ; i < para->b_num ; i ++)
    {
        t->pending_bframes[i].Y[0] = T264_malloc(t->width * t->height + (t->width * t->height >> 1), CACHE_SIZE);
    }

    t->param = *para;
    t->idr_pic_id = -1;
    t->frame_id = 0;
    t->last_i_frame_id = 0;

    T264_init_cpu(t);

    if (para->enable_rc)
    {
        rc_create(t, &t->plugins[t->plug_num]);
        t->plugins[t->plug_num].proc(t, t->plugins[t->plug_num].handle, STATE_BEFORESEQ);
        t->plug_num ++;
    }
    if (para->enable_stat)
    {
        stat_create(t, &t->plugins[t->plug_num]);
        t->plugins[t->plug_num].proc(t, t->plugins[t->plug_num].handle, STATE_BEFORESEQ);
        t->plug_num ++;
    }
    if (t->flags & USE_EXTRASUBPELSEARCH)
        t->subpel_pts = 8;
    else
        t->subpel_pts = 4;

    return t;
}

void
T264_close(T264_t* t)
{
    int32_t i;

    for(i = 0 ; i < t->param.ref_num + 1 ; i ++)
    {
        T264_free(t->refn[i].Y[0] - (EDGED_HEIGHT * t->edged_width + EDGED_WIDTH));
        T264_free(t->refn[i].mb);
        T264_free(t->refn[i].Y[1] - (EDGED_HEIGHT * t->edged_width + EDGED_WIDTH));
    }

    for(i = 0 ; i < t->plug_num ; i ++)
    {
        t->plugins[i].close(t, &t->plugins[i]);
    }

    for(i = 0 ; i < t->param.b_num ; i ++)
    {
        T264_free(t->pending_bframes[i].Y[0]);
    }
    T264_free(t->bs_buf);
    T264_free(t->bs);

    T264_free(t);	
}

static void __inline
T264_encode_frame(T264_t* t)
{
    int32_t i, j;

	//for test, to be cleaned
	//extern int frame_cabac, mb_cabac, slice_type_cabac;
    T264_load_ref(t);

    slice_header_init(t, &t->slice);
    slice_header_write(t, &t->slice);

    t->header_bits = eg_len(t->bs) * 8;
    t->sad_all = 0;
	//for CABAC
	if( t->param.cabac )
	{
		/* alignment needed */
		BitstreamPadOne( t->bs );

		/* init cabac */
		T264_cabac_context_init( &t->cabac, t->slice_type, t->ps.pic_init_qp_minus26+26+t->slice.slice_qp_delta, t->slice.cabac_init_idc );
		T264_cabac_encode_init ( &t->cabac, t->bs );
	}
    for(i = 0 ; i < t->mb_height ; i ++)
    {
        for(j = 0 ; j < t->mb_width ; j ++)
        {
            T264_mb_load_context(t, i, j);

            T264_mb_mode_decision(t);

            T264_mb_encode(t);

            T264_mb_encode_post(t);

			if(t->param.cabac)
			{
				T264_mb_save_context(t);
			}
			//SKIP
            if(t->mb.mb_mode == P_SKIP || t->mb.mb_mode == B_SKIP)
            {
				if( t->param.cabac )
				{
					if( t->mb.mb_xy > 0 )
					{
						/* not end_of_slice_flag */
						T264_cabac_encode_terminal( &t->cabac, 0 );
					}

					T264_cabac_mb_skip( t, 1 );
					//for CABAC, set MVD to zero
					memset(&(t->mb.context->mvd[0][0][0]), 0, sizeof(t->mb.context->mvd));
				}
				else
				{
                    t->skip ++;
                }
            }
            else
            {
				if( t->param.cabac )
				{
					//for CABAC, set MVD to zero
					memset(&(t->mb.mvd[0][0][0]), 0, sizeof(t->mb.mvd));
					if( t->mb.mb_xy > 0 )
					{
						/* not end_of_slice_flag */
						T264_cabac_encode_terminal( &t->cabac, 0 );
					}
					if( t->slice_type != SLICE_I )
					{
						T264_cabac_mb_skip( t, 0 );
					}
					T264_macroblock_write_cabac( t, t->bs );
					memcpy(&(t->mb.context->mvd[0][0][0]), &(t->mb.mvd[0][0][0]), sizeof(t->mb.context->mvd));
				}
                else
                {
                    T264_macroblock_write_cavlc(t);				
                }
            }
			if(!t->param.cabac)
			{
                T264_mb_save_context(t);
			}
            t->sad_all += t->mb.sad;
        }
    }

	if( t->param.cabac )
	{
		/* end of slice */
		T264_cabac_encode_terminal( &t->cabac, 1 );
	}
	else if (t->skip > 0)
    {
        eg_write_ue(t->bs, t->skip);
        t->skip = 0;
    }

	if( t->param.cabac )
	{
		int i_cabac_word;
		int pos;
		T264_cabac_encode_flush( &t->cabac );
		/* TODO cabac stuffing things (p209) */
		pos = BitstreamPos( t->bs)/8;
		i_cabac_word = (((3 * t->cabac.i_sym_cnt - 3 * 96 * t->mb_width * t->mb_height)/32) - pos)/3;

		while( i_cabac_word > 0 )
		{
			BitstreamPutBits( t->bs, 16, 0x0000 );
			i_cabac_word--;
		}
	}
	
    /* update current pic */
    if (t->slice_type != SLICE_B)
    {
        T264_save_ref(t);
        t->frame_num = (t->frame_num + 1) % (1 << (t->ss.log2_max_frame_num_minus4 + 4));
        t->poc += 2;
    }

    rbsp_trailing_bits(t);
    eg_flush(t->bs);
	//for CABAC
	if( t->param.cabac )
	{
		T264_cabac_model_update( &t->cabac, t->slice_type, t->ps.pic_init_qp_minus26+26 + t->slice.slice_qp_delta );
	}
}

static int32_t __inline
T264_flush_frames(T264_t* t, uint8_t* dst, int32_t dst_size)
{
    int32_t i, j = 0;
    int32_t nal_pos[4];     // remember each nal start pos
    int32_t nal_num;
    int32_t len = 0;
    int32_t old_poc = t->poc;

    // b frames
    while (t->pending_bframes_num)
    {
        t->pending_bframes_num --;
        t->poc = t->pending_bframes[j].poc;

        eg_init(t->bs, t->bs_buf, dst_size);
        T264_init_frame(t, t->pending_bframes[j].Y[0], &t->cur, t->poc);
        nal_num = 0;

        nal_pos[nal_num ++] = eg_len(t->bs);
        nal_unit_init(&t->nal, 0, NAL_SLICE_NOPART);
        nal_unit_write(t, &t->nal);

        t->slice_type        = SLICE_B;

        for(i = 0 ; i < t->plug_num ; i ++)
        {
            t->plugins[i].proc(t, t->plugins[i].handle, STATE_BEFOREPIC);
        }

		T264_encode_frame(t);
        nal_pos[nal_num ++] = eg_len(t->bs);

        len += write_dst(t->bs_buf, nal_pos, nal_num, dst + len, dst_size - len);

        t->emms();
 		
        t->frame_bits = len * 8;		

        for(i = 0 ; i < t->plug_num ; i ++)
        {
            t->plugins[i].proc(t, t->plugins[i].handle, STATE_AFTERPIC);
        }

        j ++;
    }

    t->poc = old_poc;

    return len;
}

int32_t
T264_encode(T264_t* t, uint8_t* src, uint8_t* dst, int32_t dst_size)
{
    int32_t i;
    int32_t nal_pos[5];     // remember each nal start pos
    int32_t nal_num = 0;
    int32_t len;

    eg_init(t->bs, t->bs_buf, dst_size);
    T264_init_frame(t, src, &t->cur, t->poc);

    t->slice_type = decision_slice_type(t);
    switch (t->slice_type)
    {
    case SLICE_P:
        nal_pos[nal_num ++] = eg_len(t->bs);
        nal_unit_init(&t->nal, 1, NAL_SLICE_NOPART);
        nal_unit_write(t, &t->nal);

        break;
    case SLICE_B:
        // buffer b frame
        T264_pending_bframe(t, src, t->poc);
        t->poc += 2;
        t->frame_id ++;
        t->frame_no ++;

        return 0;
    case SLICE_I:
        nal_pos[nal_num ++] = eg_len(t->bs);
        nal_unit_init(&t->nal, 1, NAL_SLICE_NOPART);
        nal_unit_write(t, &t->nal);

        /* reset when begin a new gop */
        t->frame_no = 0;
        t->last_i_frame_id = t->frame_id;

        for(i = 0 ; i < t->plug_num ; i ++)
        {
            t->plugins[i].proc(t, t->plugins[i].handle, STATE_BEFOREGOP);
        }

        break;
    case SLICE_IDR:
        // fixed slice type
        t->slice_type = SLICE_I;
        if (t->frame_id == 0)
        {
            /* only write once */
            T264_custom_set_t set;
            nal_pos[nal_num ++] = eg_len(t->bs);
            nal_unit_init(&t->nal, 1, NAL_CUSTOM_SET);
            nal_unit_write(t, &t->nal);
            custom_set_init(t, &set);
            custom_set_write(t, &set);
        }
        // first flush b frames
        len = T264_flush_frames(t, dst, dst_size);
        dst += len;
        dst_size -= len;
        len = 0;

        nal_pos[nal_num ++] = eg_len(t->bs);
        nal_unit_init(&t->nal, 1, NAL_SEQ_SET);
        nal_unit_write(t, &t->nal);

        seq_set_init(t, &t->ss);
        seq_set_write(t, &t->ss);

        nal_pos[nal_num ++] = eg_len(t->bs);

        nal_unit_init(&t->nal, 1, NAL_PIC_SET);
        nal_unit_write(t, &t->nal);

        pic_set_init(t, &t->ps);
        pic_set_write(t, &t->ps);

        nal_pos[nal_num ++] = eg_len(t->bs);

        nal_unit_init(&t->nal, 1, NAL_SLICE_IDR);
        nal_unit_write(t, &t->nal);

        t->idr_pic_id = (t->idr_pic_id + 1) % 65535;
        t->frame_num = 0;
        t->last_i_frame_id = t->frame_id;

        /* reset when begin a new gop */
        t->frame_no = 0;
        t->poc = 0;

        T264_reset_ref(t);

        for(i = 0 ; i < t->plug_num ; i ++)
        {
            t->plugins[i].proc(t, t->plugins[i].handle, STATE_BEFOREGOP);
        }
        break;
    default:
        assert(0);
        break;
    }

    t->frame_id ++;
    t->frame_no ++;

    for(i = 0 ; i < t->plug_num ; i ++)
    {
        t->plugins[i].proc(t, t->plugins[i].handle, STATE_BEFOREPIC);
    }

    // i or p frame
    T264_encode_frame(t);
    nal_pos[nal_num ++] = eg_len(t->bs);

    len = write_dst(t->bs_buf, nal_pos, nal_num, dst, dst_size);

    t->emms();

    t->frame_bits = len * 8;
    for(i = 0 ; i < t->plug_num ; i ++)
    {
        t->plugins[i].proc(t, t->plugins[i].handle, STATE_AFTERPIC);
    }

    len += T264_flush_frames(t, dst + len, dst_size - len);

    return len;
}