HTM-16.2代码（11）——xTZSearch

TZSearch的基本流程：

1、搜索预测得到的mv所指向的点：中值预测mv，当前PU的左，上及右上PU的mv，还有零运动矢量（0,0）；

2、在步骤1中找到匹配误差最小的点作为接下来搜索的起始点；

3、步长从1开始，以2的指数递增，进行8点钻石搜索，该步骤中可以设置搜索的最大次数（以某个步长遍历一遍就算1次）；

4、如果步骤3搜索得到的最佳步长为1，则需要以该最佳点为起点做1次两点钻石搜索，因为前面8点搜索的时候，这个最佳点的8个邻点会有两个没有搜索到；

5、如果步骤3搜索得到的最佳步长大于某个阈值（iRaster），则以步骤2得到的点作为起点，做步长为iRaster的光栅扫描（即在运动搜索的范围内遍历所有点）；

6、 最后，在经过前面1~5歩之后，以得到的最佳点为起点，再次重复步骤3和4；

7、保存最佳mv和SAD。

Void TEncSearch::xTZSearch( const TComDataCU* const pcCU,
const TComPattern* const pcPatternKey,
const Pel* const         piRefY,
const Int                iRefStride,
const TComMv* const      pcMvSrchRngLT,
const TComMv* const      pcMvSrchRngRB,
TComMv      &rcMv,
Distortion  &ruiSAD,
const TComMv* const      pIntegerMv2Nx2NPred,
const Bool               bExtendedSettings)
{
const Bool bUseAdaptiveRaster                      = bExtendedSettings;
const Int  iRaster                                 = 5;
const Bool bTestOtherPredictedMV                   = bExtendedSettings;
const Bool bTestZeroVector                         = true;
const Bool bTestZeroVectorStart                    = bExtendedSettings;
const Bool bTestZeroVectorStop                     = false;
const Bool bFirstSearchDiamond                     = true;  // 1 = xTZ8PointDiamondSearch   0 = xTZ8PointSquareSearch
const Bool bFirstCornersForDiamondDist1            = bExtendedSettings;
const Bool bFirstSearchStop                        = m_pcEncCfg->getFastMEAssumingSmootherMVEnabled();
const UInt uiFirstSearchRounds                     = 3;     // first search stop X rounds after best match (must be >=1)
const Bool bEnableRasterSearch                     = true;
const Bool bAlwaysRasterSearch                     = bExtendedSettings;  // true: BETTER but factor 2 slower
const Bool bRasterRefinementEnable                 = false; // enable either raster refinement or star refinement
const Bool bRasterRefinementDiamond                = false; // 1 = xTZ8PointDiamondSearch   0 = xTZ8PointSquareSearch
const Bool bRasterRefinementCornersForDiamondDist1 = bExtendedSettings;
const Bool bStarRefinementEnable                   = true;  // enable either star refinement or raster refinement
const Bool bStarRefinementDiamond                  = true;  // 1 = xTZ8PointDiamondSearch   0 = xTZ8PointSquareSearch
const Bool bStarRefinementCornersForDiamondDist1   = bExtendedSettings;
const Bool bStarRefinementStop                     = false;
const UInt uiStarRefinementRounds                  = 2;  // star refinement stop X rounds after best match (must be >=1)
const Bool bNewZeroNeighbourhoodTest               = bExtendedSettings;

UInt uiSearchRange = m_iSearchRange;
pcCU->clipMv( rcMv );
#if NH_3D_INTEGER_MV_DEPTH
if( ! pcCU->getSlice()->getIsDepth() )
#endif
#if ME_ENABLE_ROUNDING_OF_MVS
rcMv.divideByPowerOf2(2);
#else
rcMv >>= 2;
#endif

// init TZSearchStruct
IntTZSearchStruct cStruct;
cStruct.iYStride    = iRefStride;
cStruct.piRefY      = piRefY;
cStruct.uiBestSad   = MAX_UINT;

// set rcMv (Median predictor) as start point and as best point
xTZSearchHelp( pcPatternKey, cStruct, rcMv.getHor(), rcMv.getVer(), 0, 0 );//!< 中值预测

// test whether one of PRED_A, PRED_B, PRED_C MV is better start point than Median predictor
if ( bTestOtherPredictedMV )
{
for ( UInt index = 0; index < NUM_MV_PREDICTORS; index++ )
{
TComMv cMv = m_acMvPredictors[index];
pcCU->clipMv( cMv );
#if NH_3D_INTEGER_MV_DEPTH
if( ! pcCU->getSlice()->getIsDepth() )
{
#endif
#if ME_ENABLE_ROUNDING_OF_MVS
cMv.divideByPowerOf2(2);
#else
cMv >>= 2;
#endif

#if NH_3D_INTEGER_MV_DEPTH
}
#endif
if (cMv != rcMv && (cMv.getHor() != cStruct.iBestX && cMv.getVer() != cStruct.iBestY))
{
// only test cMV if not obviously previously tested.
xTZSearchHelp( pcPatternKey, cStruct, cMv.getHor(), cMv.getVer(), 0, 0 ); //A, B, C相邻PU的mv
}
}
}

// test whether zero Mv is better start point than Median predictor
if ( bTestZeroVector )
{
if ((rcMv.getHor() != 0 || rcMv.getVer() != 0) &&
(0 != cStruct.iBestX || 0 != cStruct.iBestY))
{
// only test 0-vector if not obviously previously tested.
xTZSearchHelp( pcPatternKey, cStruct, 0, 0, 0, 0 );//!< 零mv
}
}

Int   iSrchRngHorLeft   = pcMvSrchRngLT->getHor();
Int   iSrchRngHorRight  = pcMvSrchRngRB->getHor();
Int   iSrchRngVerTop    = pcMvSrchRngLT->getVer();
Int   iSrchRngVerBottom = pcMvSrchRngRB->getVer();

if (pIntegerMv2Nx2NPred != 0)
{
TComMv integerMv2Nx2NPred = *pIntegerMv2Nx2NPred;
integerMv2Nx2NPred <<= 2;
pcCU->clipMv( integerMv2Nx2NPred );
#if ME_ENABLE_ROUNDING_OF_MVS
integerMv2Nx2NPred.divideByPowerOf2(2);
#else
integerMv2Nx2NPred >>= 2;
#endif
if ((rcMv != integerMv2Nx2NPred) &&
(integerMv2Nx2NPred.getHor() != cStruct.iBestX || integerMv2Nx2NPred.getVer() != cStruct.iBestY))
{
// only test integerMv2Nx2NPred if not obviously previously tested.
xTZSearchHelp(pcPatternKey, cStruct, integerMv2Nx2NPred.getHor(), integerMv2Nx2NPred.getVer(), 0, 0);
}

// reset search range
TComMv cMvSrchRngLT;
TComMv cMvSrchRngRB;
Int iSrchRng = m_iSearchRange;
TComMv currBestMv(cStruct.iBestX, cStruct.iBestY );
currBestMv <<= 2;
xSetSearchRange( pcCU, currBestMv, iSrchRng, cMvSrchRngLT, cMvSrchRngRB );
iSrchRngHorLeft   = cMvSrchRngLT.getHor();
iSrchRngHorRight  = cMvSrchRngRB.getHor();
iSrchRngVerTop    = cMvSrchRngLT.getVer();
iSrchRngVerBottom = cMvSrchRngRB.getVer();
}

// start search
//从以前面几个mv作为搜索起点得到的最好的位置开始进行接下来的搜索
Int  iDist = 0;
Int  iStartX = cStruct.iBestX;
Int  iStartY = cStruct.iBestY;

const Bool bBestCandidateZero = (cStruct.iBestX == 0) && (cStruct.iBestY == 0);

// first search around best position up to now.
// The following works as a "subsampled/log" window search around the best candidate
for ( iDist = 1; iDist <= (Int)uiSearchRange; iDist*=2 )//!< 以2的幂次逐步扩大搜索步长
{
if ( bFirstSearchDiamond == 1 )
{
xTZ8PointDiamondSearch ( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB, iStartX, iStartY, iDist, bFirstCornersForDiamondDist1 );
}
else
{
xTZ8PointSquareSearch  ( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB, iStartX, iStartY, iDist );
}

if ( bFirstSearchStop && ( cStruct.uiBestRound >= uiFirstSearchRounds ) ) // stop criterion
{
break;
}
}

if (!bNewZeroNeighbourhoodTest)
{
// test whether zero Mv is a better start point than Median predictor
if ( bTestZeroVectorStart && ((cStruct.iBestX != 0) || (cStruct.iBestY != 0)) )
{
xTZSearchHelp( pcPatternKey, cStruct, 0, 0, 0, 0 );
if ( (cStruct.iBestX == 0) && (cStruct.iBestY == 0) )
{
// test its neighborhood
for ( iDist = 1; iDist <= (Int)uiSearchRange; iDist*=2 )
{
xTZ8PointDiamondSearch( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB, 0, 0, iDist, false );
if ( bTestZeroVectorStop && (cStruct.uiBestRound > 0) ) // stop criterion
{
break;
}
}
}
}
}
else
{
// Test also zero neighbourhood but with half the range
// It was reported that the original (above) search scheme using bTestZeroVectorStart did not
// make sense since one would have already checked the zero candidate earlier
// and thus the conditions for that test would have not been satisfied
if (bTestZeroVectorStart == true && bBestCandidateZero != true)
{
for ( iDist = 1; iDist <= ((Int)uiSearchRange >> 1); iDist*=2 )
{
xTZ8PointDiamondSearch( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB, 0, 0, iDist, false );
if ( bTestZeroVectorStop && (cStruct.uiBestRound > 2) ) // stop criterion
{
break;
}
}
}
}

// calculate only 2 missing points instead 8 points if cStruct.uiBestDistance == 1
if ( cStruct.uiBestDistance == 1 )//!< 当最佳搜索步长等于1时，补充搜索前面8点钻石扫描遗漏的两点
{
cStruct.uiBestDistance = 0;
xTZ2PointSearch( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB );
}

// raster search if distance is too big
////!< 当前面搜索得到的最佳步长过大时，改用光栅搜索法，步长定为iRaster，搜索范围为设定的运动估计范围
if (bUseAdaptiveRaster)
{
int iWindowSize = iRaster;
Int   iSrchRngRasterLeft   = iSrchRngHorLeft;
Int   iSrchRngRasterRight  = iSrchRngHorRight;
Int   iSrchRngRasterTop    = iSrchRngVerTop;
Int   iSrchRngRasterBottom = iSrchRngVerBottom;

if (!(bEnableRasterSearch && ( ((Int)(cStruct.uiBestDistance) > iRaster))))
{
iWindowSize ++;
iSrchRngRasterLeft /= 2;
iSrchRngRasterRight /= 2;
iSrchRngRasterTop /= 2;
iSrchRngRasterBottom /= 2;
}
cStruct.uiBestDistance = iWindowSize;
for ( iStartY = iSrchRngRasterTop; iStartY <= iSrchRngRasterBottom; iStartY += iWindowSize )
{
for ( iStartX = iSrchRngRasterLeft; iStartX <= iSrchRngRasterRight; iStartX += iWindowSize )
{
xTZSearchHelp( pcPatternKey, cStruct, iStartX, iStartY, 0, iWindowSize );
}
}
}
else
{
if ( bEnableRasterSearch && ( ((Int)(cStruct.uiBestDistance) > iRaster) || bAlwaysRasterSearch ) )
{
cStruct.uiBestDistance = iRaster;
for ( iStartY = iSrchRngVerTop; iStartY <= iSrchRngVerBottom; iStartY += iRaster )
{
for ( iStartX = iSrchRngHorLeft; iStartX <= iSrchRngHorRight; iStartX += iRaster )
{
xTZSearchHelp( pcPatternKey, cStruct, iStartX, iStartY, 0, iRaster );
}
}
}
}

// raster refinement
if ( bRasterRefinementEnable && cStruct.uiBestDistance > 0 )
{
while ( cStruct.uiBestDistance > 0 )
{
iStartX = cStruct.iBestX;
iStartY = cStruct.iBestY;
if ( cStruct.uiBestDistance > 1 )
{
iDist = cStruct.uiBestDistance >>= 1;
if ( bRasterRefinementDiamond == 1 )
{
xTZ8PointDiamondSearch ( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB, iStartX, iStartY, iDist, bRasterRefinementCornersForDiamondDist1 );
}
else
{
xTZ8PointSquareSearch  ( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB, iStartX, iStartY, iDist );
}
}

// calculate only 2 missing points instead 8 points if cStruct.uiBestDistance == 1
if ( cStruct.uiBestDistance == 1 )
{
cStruct.uiBestDistance = 0;
if ( cStruct.ucPointNr != 0 )
{
xTZ2PointSearch( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB );
}
}
}
}

// star refinement
//!< 在经过了上面几个步骤的搜索后，从最佳点开始进行第2次的8点钻石扫描以及利用两点扫描对遗漏点进行补充
if ( bStarRefinementEnable && cStruct.uiBestDistance > 0 )
{
while ( cStruct.uiBestDistance > 0 )
{
iStartX = cStruct.iBestX;
iStartY = cStruct.iBestY;
cStruct.uiBestDistance = 0;
cStruct.ucPointNr = 0;
for ( iDist = 1; iDist < (Int)uiSearchRange + 1; iDist*=2 )
{
if ( bStarRefinementDiamond == 1 )
{
xTZ8PointDiamondSearch ( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB, iStartX, iStartY, iDist, bStarRefinementCornersForDiamondDist1 );
}
else
{
xTZ8PointSquareSearch  ( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB, iStartX, iStartY, iDist );
}
if ( bStarRefinementStop && (cStruct.uiBestRound >= uiStarRefinementRounds) ) // stop criterion
{
break;
}
}

// calculate only 2 missing points instead 8 points if cStrukt.uiBestDistance == 1
if ( cStruct.uiBestDistance == 1 )
{
cStruct.uiBestDistance = 0;
if ( cStruct.ucPointNr != 0 )
{
xTZ2PointSearch( pcPatternKey, cStruct, pcMvSrchRngLT, pcMvSrchRngRB );
}
}
}
}

// write out best match   获得最佳匹配结果，mv和SAD
rcMv.set( cStruct.iBestX, cStruct.iBestY );
ruiSAD = cStruct.uiBestSad - m_pcRdCost->getCostOfVectorWithPredictor( cStruct.iBestX, cStruct.iBestY );
}

xTZSearchHelp是当中最为重要的子函数之一。它实现最基本的功能：根据输入的搜索点坐标，参考图像首地址，原始图像首地址，以及当前PU大小等相关信息，计算出SAD，并与之前保存的最佳值进行比较，更新到目前为止的最佳值相关参数，如uiBestSad，搜索点坐标，搜索步长等。其他的函数如xTZ8PointSearch等搜索函数，最终都是调用xTZSearchHelp进行误差匹配的。因此，我们有必要了解xTZSearchHelp这个函数：

__inline Void TEncSearch::xTZSearchHelp( const TComPattern* const pcPatternKey, IntTZSearchStruct& rcStruct, const Int iSearchX, const Int iSearchY, const UChar ucPointNr, const UInt uiDistance )
{
Distortion  uiSad = 0;

const Pel* const  piRefSrch = rcStruct.piRefY + iSearchY * rcStruct.iYStride + iSearchX;//参考图像Y分量的起始地址

#if NH_3D_IC
m_cDistParam.bUseIC = pcPatternKey->getICFlag();
#endif
#if NH_3D_SDC_INTER
m_cDistParam.bUseSDCMRSAD = pcPatternKey->getSDCMRSADFlag();
#endif
//-- jclee for using the SAD function pointer
m_pcRdCost->setDistParam( pcPatternKey, piRefSrch, rcStruct.iYStride,  m_cDistParam );//该函数主要职能是设置计算SAD的函数指针

setDistParamComp(COMPONENT_Y);

// distortion
m_cDistParam.bitDepth = pcPatternKey->getBitDepthY();
m_cDistParam.m_maximumDistortionForEarlyExit = rcStruct.uiBestSad;

if((m_pcEncCfg->getRestrictMESampling() == false) && m_pcEncCfg->getMotionEstimationSearchMethod() == MESEARCH_SELECTIVE)
{
Int isubShift = 0;
// motion cost
Distortion uiBitCost = m_pcRdCost->getCostOfVectorWithPredictor( iSearchX, iSearchY );

// Skip search if bit cost is already larger than best SAD
if (uiBitCost < rcStruct.uiBestSad)
{
if ( m_cDistParam.iRows > 32 )
{
m_cDistParam.iSubShift = 4;
}
else if ( m_cDistParam.iRows > 16 )
{
m_cDistParam.iSubShift = 3;
}
else if ( m_cDistParam.iRows > 8 )
{
m_cDistParam.iSubShift = 2;
}
else
{
m_cDistParam.iSubShift = 1;
}

Distortion uiTempSad = m_cDistParam.DistFunc( &m_cDistParam );
if((uiTempSad + uiBitCost) < rcStruct.uiBestSad)
{
uiSad += uiTempSad >>  m_cDistParam.iSubShift;
while(m_cDistParam.iSubShift > 0)
{
isubShift         = m_cDistParam.iSubShift -1;
m_cDistParam.pOrg = pcPatternKey->getROIY() + (pcPatternKey->getPatternLStride() << isubShift);
m_cDistParam.pCur = piRefSrch + (rcStruct.iYStride << isubShift);
uiTempSad = m_cDistParam.DistFunc( &m_cDistParam );
uiSad += uiTempSad >>  m_cDistParam.iSubShift;
if(((uiSad << isubShift) + uiBitCost) > rcStruct.uiBestSad)
{
break;
}

m_cDistParam.iSubShift--;
}

if(m_cDistParam.iSubShift == 0)
{
uiSad += uiBitCost;
if( uiSad < rcStruct.uiBestSad )
{
rcStruct.uiBestSad      = uiSad;
rcStruct.iBestX         = iSearchX;
rcStruct.iBestY         = iSearchY;
rcStruct.uiBestDistance = uiDistance;
rcStruct.uiBestRound    = 0;
rcStruct.ucPointNr      = ucPointNr;
m_cDistParam.m_maximumDistortionForEarlyExit = uiSad;
}
}
}
}
}
else
{
// fast encoder decision: use subsampled SAD when rows > 8 for integer ME
if ( m_pcEncCfg->getFastInterSearchMode()==FASTINTERSEARCH_MODE1 || m_pcEncCfg->getFastInterSearchMode()==FASTINTERSEARCH_MODE3 )
{
if ( m_cDistParam.iRows > 8 )
{
m_cDistParam.iSubShift = 1;
}
}

uiSad = m_cDistParam.DistFunc( &m_cDistParam );

// only add motion cost if uiSad is smaller than best. Otherwise pointless
// to add motion cost.
if( uiSad < rcStruct.uiBestSad )
{
// motion cost
uiSad += m_pcRdCost->getCostOfVectorWithPredictor( iSearchX, iSearchY );//!< 考虑上mv本身带来的开销

if( uiSad < rcStruct.uiBestSad )//!< 更新最佳值
{
rcStruct.uiBestSad      = uiSad;//!< SAD
rcStruct.iBestX         = iSearchX;//!< mv_x
rcStruct.iBestY         = iSearchY;
rcStruct.uiBestDistance = uiDistance;//!< 搜索步长
rcStruct.uiBestRound    = 0;//!< 搜索次数
rcStruct.ucPointNr      = ucPointNr;//!< 搜索点序号
m_cDistParam.m_maximumDistortionForEarlyExit = uiSad;
}
}
}
}

// Setting the Distortion Parameter for Inter (ME)
Void TComRdCost::setDistParam( const TComPattern* const pcPatternKey, const Pel* piRefY, Int iRefStride, DistParam& rcDistParam )
{
// set Original & Curr Pointer / Stride
rcDistParam.pOrg = pcPatternKey->getROIY();//!< 感兴趣区即待搜索的原始图像首地址
rcDistParam.pCur = piRefY;//!< 参考图像首地址

rcDistParam.iStrideOrg = pcPatternKey->getPatternLStride();//!< 原始图像跨度
rcDistParam.iStrideCur = iRefStride;//!< 参考图像跨度

// set Block Width / Height
rcDistParam.iCols    = pcPatternKey->getROIYWidth();//!< PU宽度
rcDistParam.iRows    = pcPatternKey->getROIYHeight();
rcDistParam.DistFunc = m_afpDistortFunc[DF_SAD + g_aucConvertToBit[ rcDistParam.iCols ] + 1 ];//!< 根据PU的大小选择相应的失真计算函数
rcDistParam.m_maximumDistortionForEarlyExit = std::numeric_limits<Distortion>::max();

//!< 为非对称分区AMP预测模式提供专用的失真函数
if (rcDistParam.iCols == 12)
{
rcDistParam.DistFunc = m_afpDistortFunc[DF_SAD12];//!< TComRdCost::xGetSAD12
}
else if (rcDistParam.iCols == 24)
{
rcDistParam.DistFunc = m_afpDistortFunc[DF_SAD24];//!< TComRdCost::xGetSAD24
}
else if (rcDistParam.iCols == 48)
{
rcDistParam.DistFunc = m_afpDistortFunc[DF_SAD48];//!< TComRdCost::xGetSAD48
}

#if NH_3D_DBBP
if( m_bUseMask )
{
rcDistParam.DistFunc = TComRdCost::xGetMaskedSAD;
}
#endif
// initialize
rcDistParam.iSubShift  = 0;//!< (vertical) subsampling shift (for reducing complexity)
}